From b342bbcfda5ff75a077bd650de36109591cb314f Mon Sep 17 00:00:00 2001 From: Aditya Chatterjee Date: Thu, 15 Oct 2026 09:11:22 +0000 Subject: [PATCH 1/3] Add Nemotron-TTS megakernel backend --- .pre-commit-config.yaml | 13 +- README.md | 1 + THIRD_PARTY_NOTICES.md | 40 +- app/bench_tts.cpp | 4 + app/serve.cpp | 46 +- app/synthesize.cpp | 17 + deploy/nemotron_native/LICENSE | 7 + deploy/nemotron_native/convert_model.py | 129 + .../nemotron_native/kernels/b1024/prefill.cu | 1 + .../prefill/dedup_native_prefill_impl.inc | 267 + .../native/prefill/dedup_predictor_mk2.cuh | 4600 +++++++++++++ .../prefill/native/prefill/mk2_abi.cuh | 82 + .../native/prefill/native_prefill_abi.cuh | 33 + .../regions/prefill/native/prefill/prefill.cu | 18 + .../kernels/b1024/regions/prefill/prefill.cu | 20 + .../regions/stage0/native/stage0/mk2_abi.cuh | 85 + .../native/stage0/native_prefill_abi.cuh | 33 + .../native/stage0/native_prefill_impl.inc | 198 + .../stage0/native/stage0/predictor_mk2.cu | 6111 +++++++++++++++++ .../kernels/b1024/regions/stage0/stage0.cu | 35 + .../nemotron_native/kernels/b1024/stage0.cu | 1 + .../nemotron_native/kernels/b128/prefill.cu | 1 + .../prefill/dedup_native_prefill_impl.inc | 267 + .../native/prefill/dedup_predictor_mk2.cuh | 4600 +++++++++++++ .../prefill/native/prefill/mk2_abi.cuh | 82 + .../native/prefill/native_prefill_abi.cuh | 33 + .../regions/prefill/native/prefill/prefill.cu | 18 + .../kernels/b128/regions/prefill/prefill.cu | 20 + .../regions/stage0/native/stage0/mk2_abi.cuh | 85 + .../native/stage0/native_prefill_abi.cuh | 33 + .../native/stage0/native_prefill_impl.inc | 198 + .../stage0/native/stage0/predictor_mk2.cu | 6107 ++++++++++++++++ .../kernels/b128/regions/stage0/stage0.cu | 35 + deploy/nemotron_native/kernels/b128/stage0.cu | 1 + .../nemotron_native/kernels/b2048/prefill.cu | 1 + .../prefill/dedup_native_prefill_impl.inc | 267 + .../native/prefill/dedup_predictor_mk2.cuh | 4600 +++++++++++++ .../prefill/native/prefill/mk2_abi.cuh | 82 + .../native/prefill/native_prefill_abi.cuh | 33 + .../regions/prefill/native/prefill/prefill.cu | 18 + .../kernels/b2048/regions/prefill/prefill.cu | 20 + .../regions/stage0/native/stage0/mk2_abi.cuh | 85 + .../native/stage0/native_prefill_abi.cuh | 33 + .../native/stage0/native_prefill_impl.inc | 198 + .../stage0/native/stage0/predictor_mk2.cu | 6111 +++++++++++++++++ .../kernels/b2048/regions/stage0/stage0.cu | 35 + .../nemotron_native/kernels/b2048/stage0.cu | 1 + .../nemotron_native/kernels/b256/prefill.cu | 1 + .../prefill/dedup_native_prefill_impl.inc | 267 + .../native/prefill/dedup_predictor_mk2.cuh | 4600 +++++++++++++ .../prefill/native/prefill/mk2_abi.cuh | 82 + .../native/prefill/native_prefill_abi.cuh | 33 + .../regions/prefill/native/prefill/prefill.cu | 18 + .../kernels/b256/regions/prefill/prefill.cu | 20 + .../regions/stage0/native/stage0/mk2_abi.cuh | 85 + .../native/stage0/native_prefill_abi.cuh | 33 + .../native/stage0/native_prefill_impl.inc | 198 + .../stage0/native/stage0/predictor_mk2.cu | 6107 ++++++++++++++++ .../kernels/b256/regions/stage0/stage0.cu | 35 + deploy/nemotron_native/kernels/b256/stage0.cu | 1 + deploy/nemotron_native/kernels/b32/codec.cu | 4 + .../b32/native/codec/kernel_kf_codec.h | 13 + .../kernels/b32/native/codec/legacy_codec.cuh | 1390 ++++ .../kernels/b32/native/codec/native_stage1.cu | 329 + .../kernels/b32/native/codec/native_stage1.h | 26 + .../kernels/b32/native/codec/repack3.inc | 66 + .../kernels/b32/native/codec/tail_kernel.inc | 492 ++ .../kernels/b32/native/codec/wg2_kernel.inc | 280 + .../kernels/b32/native/codec/wg3_kernel.inc | 177 + .../b32/native/codec/wg4_step_widths.inc | 459 ++ .../b32/native/codec/wg5_step_widths.inc | 469 ++ .../kernels/b32/native/codec/wg_kernel.inc | 182 + .../kernels/b32/native/codec/wgmma_macro.inc | 102 + .../prefill/dedup_native_prefill_impl.inc | 249 + .../native/prefill/dedup_predictor_mk2.cuh | 4600 +++++++++++++ .../kernels/b32/native/prefill/mk2_abi.cuh | 82 + .../b32/native/prefill/native_prefill_abi.cuh | 33 + .../kernels/b32/native/prefill/prefill.cu | 18 + .../kernels/b32/native/stage0/mk2_abi.cuh | 91 + .../b32/native/stage0/native_prefill_abi.cuh | 33 + .../b32/native/stage0/native_prefill_impl.inc | 198 + .../b32/native/stage0/predictor_mk2.cu | 5472 +++++++++++++++ .../b32/native/stage0/sparse_moe_impl.cuh | 388 ++ deploy/nemotron_native/kernels/b32/prefill.cu | 20 + deploy/nemotron_native/kernels/b32/stage0.cu | 35 + deploy/nemotron_native/kernels/b4096/codec.cu | 1 + .../nemotron_native/kernels/b4096/prefill.cu | 1 + .../prefill/dedup_native_prefill_impl.inc | 267 + .../native/prefill/dedup_predictor_mk2.cuh | 4600 +++++++++++++ .../prefill/native/prefill/mk2_abi.cuh | 82 + .../native/prefill/native_prefill_abi.cuh | 33 + .../regions/prefill/native/prefill/prefill.cu | 18 + .../kernels/b4096/regions/prefill/prefill.cu | 20 + .../regions/stage0/native/stage0/mk2_abi.cuh | 85 + .../native/stage0/native_prefill_abi.cuh | 33 + .../native/stage0/native_prefill_impl.inc | 198 + .../stage0/native/stage0/predictor_mk2.cu | 6107 ++++++++++++++++ .../kernels/b4096/regions/stage0/stage0.cu | 35 + .../kernels/b4096/regions/stage1/codec.cu | 4 + .../stage1/native/codec/kernel_kf_codec.h | 13 + .../stage1/native/codec/legacy_codec.cuh | 1387 ++++ .../stage1/native/codec/native_stage1.cu | 357 + .../stage1/native/codec/native_stage1.h | 26 + .../regions/stage1/native/codec/repack3.inc | 66 + .../stage1/native/codec/tail_kernel.inc | 492 ++ .../stage1/native/codec/wg2_kernel.inc | 280 + .../stage1/native/codec/wg3_kernel.inc | 177 + .../stage1/native/codec/wg4_step_widths.inc | 459 ++ .../stage1/native/codec/wg5_step_widths.inc | 400 ++ .../regions/stage1/native/codec/wg_kernel.inc | 182 + .../stage1/native/codec/wgmma_macro.inc | 102 + .../nemotron_native/kernels/b4096/stage0.cu | 1 + deploy/nemotron_native/kernels/b512/codec.cu | 4 + .../b512/native/codec/kernel_kf_codec.h | 13 + .../b512/native/codec/legacy_codec.cuh | 1389 ++++ .../b512/native/codec/native_stage1.cu | 359 + .../kernels/b512/native/codec/native_stage1.h | 26 + .../kernels/b512/native/codec/repack3.inc | 66 + .../kernels/b512/native/codec/tail_kernel.inc | 502 ++ .../kernels/b512/native/codec/wg2_kernel.inc | 280 + .../kernels/b512/native/codec/wg3_kernel.inc | 177 + .../b512/native/codec/wg4_step_widths.inc | 468 ++ .../b512/native/codec/wg5_step_widths.inc | 528 ++ .../kernels/b512/native/codec/wg_kernel.inc | 182 + .../kernels/b512/native/codec/wgmma_macro.inc | 102 + .../prefill/dedup_native_prefill_impl.inc | 249 + .../native/prefill/dedup_predictor_mk2.cuh | 4594 +++++++++++++ .../kernels/b512/native/prefill/mk2_abi.cuh | 82 + .../native/prefill/native_prefill_abi.cuh | 33 + .../kernels/b512/native/prefill/prefill.cu | 18 + .../kernels/b512/native/stage0/mk2_abi.cuh | 85 + .../b512/native/stage0/native_prefill_abi.cuh | 33 + .../native/stage0/native_prefill_impl.inc | 198 + .../b512/native/stage0/predictor_mk2.cu | 5563 +++++++++++++++ .../nemotron_native/kernels/b512/prefill.cu | 20 + deploy/nemotron_native/kernels/b512/stage0.cu | 35 + .../kernels/codec/base/kernel_kf_codec.h | 13 + .../kernels/codec/base/legacy_codec.cuh | 1388 ++++ .../kernels/codec/base/native_stage1.cu | 309 + .../kernels/codec/base/native_stage1.h | 26 + .../kernels/codec/base/repack3.inc | 66 + .../kernels/codec/base/tail_kernel.inc | 492 ++ .../kernels/codec/base/wg2_kernel.inc | 280 + .../kernels/codec/base/wg3_kernel.inc | 177 + .../kernels/codec/base/wg4_step_widths.inc | 459 ++ .../kernels/codec/base/wg5_step_widths.inc | 400 ++ .../kernels/codec/base/wg_kernel.inc | 182 + .../kernels/codec/base/wgmma_macro.inc | 102 + .../kernels/codec/first_packet/codec.cu | 1 + .../first_packet/regions/stage1/codec.cu | 4 + .../stage1/native/codec/kernel_kf_codec.h | 13 + .../stage1/native/codec/legacy_codec.cuh | 1389 ++++ .../stage1/native/codec/native_stage1.cu | 357 + .../stage1/native/codec/native_stage1.h | 26 + .../regions/stage1/native/codec/repack3.inc | 66 + .../stage1/native/codec/tail_kernel.inc | 502 ++ .../stage1/native/codec/wg2_kernel.inc | 280 + .../stage1/native/codec/wg3_kernel.inc | 177 + .../stage1/native/codec/wg4_step_widths.inc | 468 ++ .../stage1/native/codec/wg5_step_widths.inc | 409 ++ .../regions/stage1/native/codec/wg_kernel.inc | 182 + .../stage1/native/codec/wgmma_macro.inc | 102 + .../kernels/codec/small/kernel_kf_codec.h | 13 + .../kernels/codec/small/legacy_codec.cuh | 1388 ++++ .../kernels/codec/small/native_stage1.cu | 310 + .../kernels/codec/small/native_stage1.h | 26 + .../kernels/codec/small/repack3.inc | 66 + .../kernels/codec/small/tail_kernel.inc | 492 ++ .../kernels/codec/small/wg2_kernel.inc | 280 + .../kernels/codec/small/wg3_kernel.inc | 177 + .../kernels/codec/small/wg4_step_widths.inc | 459 ++ .../kernels/codec/small/wg5_step_widths.inc | 400 ++ .../kernels/codec/small/wg_kernel.inc | 182 + .../kernels/codec/small/wgmma_macro.inc | 102 + .../kernels/contraction/bf16/mk2_abi.cuh | 248 + .../contraction/bf16/native_prefill_abi.cuh | 69 + .../contraction/bf16/native_prefill_impl.inc | 216 + .../kernels/contraction/bf16/predictor_mk2.cu | 5642 +++++++++++++++ .../contraction/bf16/sparse_moe_impl.cuh | 156 + .../kernels/contraction/lane_transfer.cu | 67 + .../kernels/contraction/mk2_abi.cuh | 248 + .../kernels/inflight/progress.cuh | 17 + .../kernels/inflight/runtime.cuh | 141 + .../prefill/b1/dedup_native_prefill_impl.inc | 237 + .../prefill/b1/dedup_predictor_mk2.cuh | 4588 +++++++++++++ .../kernels/prefill/b1/mk2_abi.cuh | 82 + .../kernels/prefill/b1/native_prefill_abi.cuh | 33 + .../kernels/prefill/b1/prefill.cu | 18 + .../prefill/b16/dedup_native_prefill_impl.inc | 237 + .../prefill/b16/dedup_predictor_mk2.cuh | 4588 +++++++++++++ .../kernels/prefill/b16/mk2_abi.cuh | 82 + .../prefill/b16/native_prefill_abi.cuh | 33 + .../kernels/prefill/b16/prefill.cu | 18 + .../prefill/b2/dedup_native_prefill_impl.inc | 237 + .../prefill/b2/dedup_predictor_mk2.cuh | 4588 +++++++++++++ .../kernels/prefill/b2/mk2_abi.cuh | 82 + .../kernels/prefill/b2/native_prefill_abi.cuh | 33 + .../kernels/prefill/b2/prefill.cu | 18 + .../prefill/b4/dedup_native_prefill_impl.inc | 237 + .../prefill/b4/dedup_predictor_mk2.cuh | 4588 +++++++++++++ .../kernels/prefill/b4/mk2_abi.cuh | 82 + .../kernels/prefill/b4/native_prefill_abi.cuh | 33 + .../kernels/prefill/b4/prefill.cu | 18 + .../prefill/b64/dedup_native_prefill_impl.inc | 237 + .../prefill/b64/dedup_predictor_mk2.cuh | 4588 +++++++++++++ .../kernels/prefill/b64/mk2_abi.cuh | 82 + .../prefill/b64/native_prefill_abi.cuh | 33 + .../kernels/prefill/b64/prefill.cu | 18 + .../prefill/b8/dedup_native_prefill_impl.inc | 237 + .../prefill/b8/dedup_predictor_mk2.cuh | 4588 +++++++++++++ .../kernels/prefill/b8/mk2_abi.cuh | 82 + .../kernels/prefill/b8/native_prefill_abi.cuh | 33 + .../kernels/prefill/b8/prefill.cu | 18 + deploy/nemotron_native/kernels/recipes.json | 2936 ++++++++ .../kernels/stage0/b1/bb_router_selection.cuh | 39 + .../kernels/stage0/b1/mk2_abi.cuh | 85 + .../kernels/stage0/b1/native_prefill_abi.cuh | 33 + .../kernels/stage0/b1/native_prefill_impl.inc | 198 + .../kernels/stage0/b1/predictor_mk2.cu | 3910 +++++++++++ .../kernels/stage0/b16/mk2_abi.cuh | 85 + .../kernels/stage0/b16/native_prefill_abi.cuh | 33 + .../stage0/b16/native_prefill_impl.inc | 198 + .../kernels/stage0/b16/predictor_mk2.cu | 4073 +++++++++++ .../kernels/stage0/b16/sparse_moe_impl.cuh | 139 + .../kernels/stage0/b64/mk2_abi.cuh | 85 + .../kernels/stage0/b64/native_prefill_abi.cuh | 33 + .../stage0/b64/native_prefill_impl.inc | 198 + .../kernels/stage0/b64/predictor_mk2.cu | 4022 +++++++++++ .../kernels/stage0/small/mk2_abi.cuh | 85 + .../stage0/small/native_prefill_abi.cuh | 33 + .../stage0/small/native_prefill_impl.inc | 198 + .../kernels/stage0/small/predictor_mk2.cu | 4074 +++++++++++ .../kernels/stage0/small/sparse_moe_impl.cuh | 139 + .../thunderkittens-v4-prefill/LICENSE | 21 + .../include/common/base_ops.cuh | 446 ++ .../include/common/base_types.cuh | 777 +++ .../include/common/common.cuh | 15 + .../include/common/multimem.cuh | 407 ++ .../include/common/util.cuh | 551 ++ .../include/kittens.cuh | 73 + .../include/ops/group/group.cuh | 117 + .../include/ops/group/memory/memory.cuh | 7 + .../tile/complex_global_to_register.cuh | 42 + .../memory/tile/complex_global_to_shared.cuh | 69 + .../tile/complex_shared_to_register.cuh | 32 + .../group/memory/tile/global_to_register.cuh | 199 + .../group/memory/tile/global_to_shared.cuh | 167 + .../memory/tile/parallel_global_to_global.cuh | 51 + .../group/memory/tile/shared_to_register.cuh | 321 + .../group/memory/tile/tensor_to_register.cuh | 493 ++ .../include/ops/group/memory/tile/tile.cuh | 20 + .../include/ops/group/memory/tile/tma.cuh | 134 + .../ops/group/memory/tile/tma_cluster.cuh | 32 + .../group/memory/vec/global_to_register.cuh | 138 + .../ops/group/memory/vec/global_to_shared.cuh | 113 + .../group/memory/vec/shared_to_register.cuh | 159 + .../include/ops/group/memory/vec/tma.cuh | 221 + .../ops/group/memory/vec/tma_cluster.cuh | 31 + .../include/ops/group/memory/vec/vec.cuh | 10 + .../include/ops/group/mma/base/64x112.impl | 334 + .../include/ops/group/mma/base/64x128.impl | 1029 +++ .../include/ops/group/mma/base/64x144.impl | 382 ++ .../include/ops/group/mma/base/64x16.impl | 190 + .../include/ops/group/mma/base/64x160.impl | 914 +++ .../include/ops/group/mma/base/64x176.impl | 430 ++ .../include/ops/group/mma/base/64x192.impl | 954 +++ .../include/ops/group/mma/base/64x208.impl | 478 ++ .../include/ops/group/mma/base/64x224.impl | 1138 +++ .../include/ops/group/mma/base/64x240.impl | 526 ++ .../include/ops/group/mma/base/64x256.impl | 1604 +++++ .../include/ops/group/mma/base/64x32.impl | 564 ++ .../include/ops/group/mma/base/64x48.impl | 238 + .../include/ops/group/mma/base/64x64.impl | 739 ++ .../include/ops/group/mma/base/64x80.impl | 286 + .../include/ops/group/mma/base/64x96.impl | 887 +++ .../include/ops/group/mma/base/base.cuh | 47 + .../include/ops/group/mma/mma.cuh | 17 + .../include/ops/group/mma/tcgen05.cuh | 236 + .../include/ops/group/mma/warp.cuh | 1057 +++ .../include/ops/group/mma/warpgroup.cuh | 1167 ++++ .../include/ops/group/register/register.cuh | 7 + .../register/tile/complex_conversions.cuh | 98 + .../ops/group/register/tile/complex_maps.cuh | 137 + .../ops/group/register/tile/conversions.cuh | 917 +++ .../include/ops/group/register/tile/maps.cuh | 890 +++ .../ops/group/register/tile/reductions.cuh | 540 ++ .../include/ops/group/register/tile/tile.cuh | 46 + .../ops/group/register/vec/conversions.cuh | 153 + .../include/ops/group/register/vec/maps.cuh | 374 + .../ops/group/register/vec/reductions.cuh | 233 + .../include/ops/group/register/vec/vec.cuh | 59 + .../include/ops/group/shared/shared.cuh | 7 + .../ops/group/shared/tile/conversions.cuh | 16 + .../include/ops/group/shared/tile/maps.cuh | 487 ++ .../ops/group/shared/tile/reductions.cuh | 372 + .../include/ops/group/shared/tile/tile.cuh | 37 + .../ops/group/shared/vec/conversions.cuh | 27 + .../include/ops/group/shared/vec/maps.cuh | 259 + .../ops/group/shared/vec/reductions.cuh | 207 + .../include/ops/group/shared/vec/vec.cuh | 37 + .../include/ops/group/util/sync.cuh | 224 + .../include/ops/group/util/tma.cuh | 90 + .../include/ops/group/util/tma_cluster.cuh | 118 + .../include/ops/group/util/util.cuh | 94 + .../include/ops/ops.cuh | 261 + .../include/ops/thread/memory/memory.cuh | 9 + .../thread/memory/tile/shared_to_tensor.cuh | 39 + .../include/ops/thread/memory/tile/tile.cuh | 13 + .../include/ops/thread/memory/tile/tma.cuh | 880 +++ .../include/ops/thread/memory/vec/tma.cuh | 418 ++ .../include/ops/thread/memory/vec/vec.cuh | 10 + .../include/ops/thread/mma/mma.cuh | 10 + .../include/ops/thread/mma/tcgen05.cuh | 985 +++ .../include/ops/thread/thread.cuh | 11 + .../include/ops/thread/util/sync.cuh | 331 + .../include/ops/thread/util/tma.cuh | 280 + .../include/ops/thread/util/util.cuh | 327 + .../include/pyutils/broker.cuh | 549 ++ .../include/pyutils/club.cuh | 124 + .../include/pyutils/parallel_tensor.cuh | 336 + .../include/pyutils/profiler.cuh | 148 + .../include/pyutils/pyutils.cuh | 94 + .../include/pyutils/torchutils.cuh | 227 + .../include/types/global/cgl.cuh | 50 + .../include/types/global/gl.cuh | 250 + .../include/types/global/global.cuh | 10 + .../include/types/global/tma.cuh | 337 + .../include/types/global/util.cuh | 180 + .../include/types/register/crt.cuh | 95 + .../include/types/register/crv.cuh | 85 + .../include/types/register/register.cuh | 14 + .../include/types/register/rt.cuh | 274 + .../include/types/register/rt_base.cuh | 124 + .../include/types/register/rt_layout.cuh | 40 + .../include/types/register/rv.cuh | 230 + .../include/types/register/rv_layout.cuh | 38 + .../include/types/shared/cst.cuh | 78 + .../include/types/shared/csv.cuh | 72 + .../include/types/shared/descriptor.cuh | 142 + .../include/types/shared/shared.cuh | 14 + .../include/types/shared/st.cuh | 658 ++ .../include/types/shared/sv.cuh | 155 + .../include/types/system/ipc.cuh | 143 + .../include/types/system/pgl.cuh | 182 + .../include/types/system/system.cuh | 14 + .../include/types/system/vmm.cuh | 176 + .../include/types/tensor/tensor.cuh | 152 + .../include/types/tensor/tt.cuh | 135 + .../include/types/types.cuh | 67 + .../third_party/thunderkittens/LICENSE | 21 + .../include/common/base_ops.cuh | 446 ++ .../include/common/base_types.cuh | 777 +++ .../thunderkittens/include/common/common.cuh | 15 + .../include/common/multimem.cuh | 407 ++ .../thunderkittens/include/common/util.cuh | 551 ++ .../thunderkittens/include/kittens.cuh | 73 + .../include/ops/group/group.cuh | 117 + .../include/ops/group/memory/memory.cuh | 7 + .../tile/complex_global_to_register.cuh | 42 + .../memory/tile/complex_global_to_shared.cuh | 69 + .../tile/complex_shared_to_register.cuh | 32 + .../group/memory/tile/global_to_register.cuh | 199 + .../group/memory/tile/global_to_shared.cuh | 167 + .../memory/tile/parallel_global_to_global.cuh | 51 + .../group/memory/tile/shared_to_register.cuh | 321 + .../group/memory/tile/tensor_to_register.cuh | 493 ++ .../include/ops/group/memory/tile/tile.cuh | 20 + .../include/ops/group/memory/tile/tma.cuh | 134 + .../ops/group/memory/tile/tma_cluster.cuh | 32 + .../group/memory/vec/global_to_register.cuh | 138 + .../ops/group/memory/vec/global_to_shared.cuh | 113 + .../group/memory/vec/shared_to_register.cuh | 159 + .../include/ops/group/memory/vec/tma.cuh | 221 + .../ops/group/memory/vec/tma_cluster.cuh | 31 + .../include/ops/group/memory/vec/vec.cuh | 10 + .../include/ops/group/mma/base/64x112.impl | 334 + .../include/ops/group/mma/base/64x128.impl | 1029 +++ .../include/ops/group/mma/base/64x144.impl | 382 ++ .../include/ops/group/mma/base/64x16.impl | 190 + .../include/ops/group/mma/base/64x160.impl | 914 +++ .../include/ops/group/mma/base/64x176.impl | 430 ++ .../include/ops/group/mma/base/64x192.impl | 954 +++ .../include/ops/group/mma/base/64x208.impl | 478 ++ .../include/ops/group/mma/base/64x224.impl | 1138 +++ .../include/ops/group/mma/base/64x240.impl | 526 ++ .../include/ops/group/mma/base/64x256.impl | 1604 +++++ .../include/ops/group/mma/base/64x32.impl | 564 ++ .../include/ops/group/mma/base/64x48.impl | 238 + .../include/ops/group/mma/base/64x64.impl | 739 ++ .../include/ops/group/mma/base/64x80.impl | 286 + .../include/ops/group/mma/base/64x96.impl | 887 +++ .../include/ops/group/mma/base/base.cuh | 47 + .../include/ops/group/mma/mma.cuh | 17 + .../include/ops/group/mma/tcgen05.cuh | 236 + .../include/ops/group/mma/warp.cuh | 1057 +++ .../include/ops/group/mma/warpgroup.cuh | 1167 ++++ .../include/ops/group/register/register.cuh | 7 + .../register/tile/complex_conversions.cuh | 98 + .../ops/group/register/tile/complex_maps.cuh | 137 + .../ops/group/register/tile/conversions.cuh | 917 +++ .../include/ops/group/register/tile/maps.cuh | 890 +++ .../ops/group/register/tile/reductions.cuh | 540 ++ .../include/ops/group/register/tile/tile.cuh | 46 + .../ops/group/register/vec/conversions.cuh | 153 + .../include/ops/group/register/vec/maps.cuh | 374 + .../ops/group/register/vec/reductions.cuh | 233 + .../include/ops/group/register/vec/vec.cuh | 59 + .../include/ops/group/shared/shared.cuh | 7 + .../ops/group/shared/tile/conversions.cuh | 16 + .../include/ops/group/shared/tile/maps.cuh | 487 ++ .../ops/group/shared/tile/reductions.cuh | 372 + .../include/ops/group/shared/tile/tile.cuh | 37 + .../ops/group/shared/vec/conversions.cuh | 27 + .../include/ops/group/shared/vec/maps.cuh | 259 + .../ops/group/shared/vec/reductions.cuh | 207 + .../include/ops/group/shared/vec/vec.cuh | 37 + .../include/ops/group/util/sync.cuh | 224 + .../include/ops/group/util/tma.cuh | 90 + .../include/ops/group/util/tma_cluster.cuh | 118 + .../include/ops/group/util/util.cuh | 94 + .../thunderkittens/include/ops/ops.cuh | 261 + .../include/ops/thread/memory/memory.cuh | 9 + .../thread/memory/tile/shared_to_tensor.cuh | 39 + .../include/ops/thread/memory/tile/tile.cuh | 13 + .../include/ops/thread/memory/tile/tma.cuh | 880 +++ .../include/ops/thread/memory/vec/tma.cuh | 418 ++ .../include/ops/thread/memory/vec/vec.cuh | 10 + .../include/ops/thread/mma/mma.cuh | 10 + .../include/ops/thread/mma/tcgen05.cuh | 985 +++ .../include/ops/thread/thread.cuh | 11 + .../include/ops/thread/util/sync.cuh | 331 + .../include/ops/thread/util/tma.cuh | 280 + .../include/ops/thread/util/util.cuh | 327 + .../thunderkittens/include/pyutils/broker.cuh | 549 ++ .../thunderkittens/include/pyutils/club.cuh | 124 + .../include/pyutils/parallel_tensor.cuh | 336 + .../include/pyutils/profiler.cuh | 148 + .../include/pyutils/pyutils.cuh | 94 + .../include/pyutils/torchutils.cuh | 227 + .../include/types/global/cgl.cuh | 50 + .../include/types/global/gl.cuh | 250 + .../include/types/global/global.cuh | 10 + .../include/types/global/tma.cuh | 337 + .../include/types/global/util.cuh | 180 + .../include/types/register/crt.cuh | 95 + .../include/types/register/crv.cuh | 85 + .../include/types/register/register.cuh | 14 + .../include/types/register/rt.cuh | 274 + .../include/types/register/rt_base.cuh | 124 + .../include/types/register/rt_layout.cuh | 40 + .../include/types/register/rv.cuh | 230 + .../include/types/register/rv_layout.cuh | 38 + .../include/types/shared/cst.cuh | 78 + .../include/types/shared/csv.cuh | 72 + .../include/types/shared/descriptor.cuh | 142 + .../include/types/shared/shared.cuh | 14 + .../include/types/shared/st.cuh | 658 ++ .../include/types/shared/sv.cuh | 155 + .../include/types/system/ipc.cuh | 143 + .../include/types/system/pgl.cuh | 182 + .../include/types/system/system.cuh | 14 + .../include/types/system/vmm.cuh | 176 + .../include/types/tensor/tensor.cuh | 152 + .../include/types/tensor/tt.cuh | 135 + .../thunderkittens/include/types/types.cuh | 67 + .../prototype/common/common.cuh | 4 + .../prototype/common/templates.cuh | 110 + .../thunderkittens/prototype/common/util.cuh | 13 + .../prototype/interpreter/interpreter.cuh | 440 ++ .../prototype/interpreter/templates.cuh | 178 + .../thunderkittens/prototype/lcf/lcf.cuh | 188 + .../prototype/lcf/templates.cuh | 121 + .../thunderkittens/prototype/lcsc/lcsc.cuh | 117 + .../thunderkittens/prototype/lcsf/lcsf.cuh | 229 + .../prototype/lcsf/templates.cuh | 141 + .../thunderkittens/prototype/prototype.cuh | 14 + deploy/nemotron_native/model.lock.json | 122 + .../profiles/playback-safe.json | 28 + deploy/nemotron_native/programs.json | 1943 ++++++ deploy/nemotron_native/serve.py | 233 + docker/Dockerfile.nemotron | 72 + docs/nemotron-integration-results.md | 227 + docs/nemotron-megakernel.md | 328 + include/nemo_speech/tts.h | 3 + native/tts-normalizer/.bazelrc | 8 + native/tts-normalizer/BUILD.bazel | 97 + native/tts-normalizer/LICENSE.legacy | 7 + native/tts-normalizer/LICENSE.riva | 7 + native/tts-normalizer/README.md | 49 + native/tts-normalizer/THIRD_PARTY_NOTICES.md | 24 + native/tts-normalizer/WORKSPACE.bazel | 89 + native/tts-normalizer/include/processor_abi.h | 129 + native/tts-normalizer/src/normalize.cc | 221 + native/tts-normalizer/src/normalize.h | 147 + native/tts-normalizer/src/speech_hints.cc | 359 + native/tts-normalizer/src/speech_hints.h | 198 + .../tts-normalizer/src/text_preprocessor.cc | 510 ++ native/tts-normalizer/src/text_preprocessor.h | 107 + .../tts-normalizer/src/text_preprocessor_c.cc | 330 + native/tts-normalizer/testdata/BUILD.bazel | 5 + native/tts-normalizer/testdata/en-US/abbr.txt | 18 + .../testdata/en-US/post_process.far | Bin 0 -> 327407 bytes .../testdata/en-US/postprocessor.ascii_proto | 5 + .../en-US/sentence_boundary_exceptions.txt | 15 + .../sparrowhawk_configuration.ascii_proto | 9 + .../testdata/en-US/tokenize_and_classify.far | Bin 0 -> 2390007 bytes .../testdata/en-US/tokenizer.ascii_proto | 5 + .../testdata/en-US/verbalize.far | Bin 0 -> 1686728 bytes .../testdata/en-US/verbalizer.ascii_proto | 5 + native/tts-normalizer/tests/normalize_test.cc | 119 + .../tts-normalizer/tests/speech_hints_test.cc | 47 + .../tests/string_matcher_test.cc | 294 + .../tests/text_preprocessor_c_test.cc | 246 + .../tests/text_preprocessor_test.cc | 236 + native/tts-normalizer/third_party/BUILD.bazel | 3 + .../tts-normalizer/third_party/BUILD.openfst | 1212 ++++ .../tts-normalizer/third_party/BUILD.pugixml | 24 + native/tts-normalizer/third_party/BUILD.re2 | 282 + .../third_party/BUILD.sparrowhawk_0_1 | 121 + .../third_party/BUILD.thrax_1_2_7 | 167 + .../third_party/sparrowhawk_normalizer.cc | 437 ++ .../third_party/sparrowhawk_rule_system.cc | 347 + .../third_party/sparrowhawk_string_matcher.h | 583 ++ python/nemotron_cpp_bridge/__init__.py | 3 + python/nemotron_cpp_bridge/__main__.py | 5 + python/nemotron_cpp_bridge/normalization.py | 344 + python/nemotron_cpp_bridge/protocol.py | 102 + python/nemotron_cpp_bridge/server.py | 581 ++ python/nemotron_cpp_bridge/slabs.py | 78 + python/nemotron_native_backend/__init__.py | 3 + python/nemotron_native_backend/backend.py | 87 + python/nemotron_native_backend/batcher.py | 300 + python/nemotron_native_backend/bundle.py | 197 + .../execution/__init__.py | 0 .../execution/chunk_boundary.py | 9 + .../execution/chunk_submission.py | 44 + .../execution/composition.py | 287 + .../execution/controller_rules.py | 121 + .../execution/engine.py | 1220 ++++ .../execution/helper_graphs.py | 280 + .../execution/inflight_engine.py | 1367 ++++ .../execution/inflight_policy.py | 69 + .../execution/inflight_reset.py | 164 + .../execution/lane_control.py | 37 + .../execution/lane_windows.py | 100 + .../execution/narrow_prefix_cache.py | 109 + .../execution/narrow_state_binding.py | 70 + .../execution/native_assets.py | 477 ++ .../execution/native_inputs.py | 190 + .../execution/native_stage1.py | 534 ++ .../execution/pcm_transfer.py | 112 + .../execution/prefix_cache.py | 84 + .../execution/prefix_snapshot.py | 246 + .../execution/program.py | 286 + .../execution/readiness.py | 265 + .../execution/request_noise.py | 48 + .../execution/research_program.py | 131 + .../execution/retained_narrow.py | 198 + .../execution/reusable_inputs.py | 129 + .../execution/reusable_runner.py | 50 + .../execution/small_width.py | 353 + .../execution/source_build.py | 154 + .../execution/split_native_runtime.py | 316 + .../execution/startup_binding.py | 24 + .../execution/state_binding.py | 64 + .../execution/token_staging.py | 102 + .../execution/wide_noise.py | 123 + python/nemotron_native_backend/hardware.py | 61 + .../nemotron_native_backend/host/__init__.py | 0 .../nemotron_native_backend/host/cpu_pools.py | 49 + .../host/diagnostic.py | 33 + .../host/exact_admission.py | 169 + .../host/first_fence.py | 130 + .../host/ready_writer.py | 118 + .../host/serialization.py | 145 + .../host/slab_audit.py | 61 + .../host/slab_cancel_v2.py | 6 + .../nemotron_native_backend/host/slab_pool.py | 115 + .../nemotron_native_backend/host/transport.py | 418 ++ .../inflight_batcher.py | 502 ++ .../initial_publication.py | 207 + python/nemotron_native_backend/lifecycle.py | 132 + python/nemotron_native_backend/manifest.py | 15 + python/nemotron_native_backend/models.py | 162 + python/nemotron_native_backend/protocol.py | 285 + python/nemotron_native_backend/r31.py | 97 + python/nemotron_native_backend/receipts.py | 83 + python/nemotron_native_backend/server.py | 354 + .../streaming_batcher.py | 534 ++ .../streaming_contract.py | 475 ++ .../streaming_models.py | 130 + .../streaming_server.py | 305 + python/nemotron_native_backend/worker.py | 632 ++ scripts/check_license_headers.py | 15 + scripts/nemotron_benchmark.py | 537 ++ scripts/nemotron_native/build_kernels.py | 300 + scripts/nemotron_native/build_package.py | 423 ++ scripts/nemotron_native/launch.py | 274 + scripts/nemotron_native/requirements.txt | 4 + .../nemotron_native/runtime_dependencies.py | 245 + scripts/nemotron_native/tests/test_bridge.py | 533 ++ .../nemotron_native/tests/test_launcher.py | 117 + server/http/CMakeLists.txt | 1 + server/http/http_server.cpp | 43 +- src/server/riva_server.cc | 5 + src/tts/c_api.cpp | 7 + src/tts/magpietts/CMakeLists.txt | 4 + src/tts/magpietts/config.cpp | 3 + src/tts/magpietts/config.h | 3 + src/tts/nemotron/native_runtime.cpp | 544 ++ src/tts/nemotron/native_runtime.h | 33 + src/tts/nemotron/pcm_conversion.h | 108 + src/tts/synthesizer.cpp | 177 +- src/tts/synthesizer.h | 8 + tests/cpp/CMakeLists.txt | 15 + tests/cpp/common/test_nemotron_pcm.cpp | 210 + tests/nemotron_http_test.py | 280 + tests/nemotron_mock_worker.py | 172 + tests/nemotron_normalization_test.py | 104 + 619 files changed, 254637 insertions(+), 38 deletions(-) create mode 100644 deploy/nemotron_native/LICENSE create mode 100644 deploy/nemotron_native/convert_model.py create mode 100644 deploy/nemotron_native/kernels/b1024/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b1024/regions/prefill/native/prefill/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b1024/regions/prefill/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b1024/regions/stage0/native/stage0/predictor_mk2.cu create mode 100644 deploy/nemotron_native/kernels/b1024/regions/stage0/stage0.cu create mode 100644 deploy/nemotron_native/kernels/b1024/stage0.cu create mode 100644 deploy/nemotron_native/kernels/b128/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b128/regions/prefill/native/prefill/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b128/regions/prefill/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b128/regions/stage0/native/stage0/predictor_mk2.cu create mode 100644 deploy/nemotron_native/kernels/b128/regions/stage0/stage0.cu create mode 100644 deploy/nemotron_native/kernels/b128/stage0.cu create mode 100644 deploy/nemotron_native/kernels/b2048/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b2048/regions/prefill/native/prefill/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b2048/regions/prefill/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b2048/regions/stage0/native/stage0/predictor_mk2.cu create mode 100644 deploy/nemotron_native/kernels/b2048/regions/stage0/stage0.cu create mode 100644 deploy/nemotron_native/kernels/b2048/stage0.cu create mode 100644 deploy/nemotron_native/kernels/b256/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b256/regions/prefill/native/prefill/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b256/regions/prefill/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b256/regions/stage0/native/stage0/predictor_mk2.cu create mode 100644 deploy/nemotron_native/kernels/b256/regions/stage0/stage0.cu create mode 100644 deploy/nemotron_native/kernels/b256/stage0.cu create mode 100644 deploy/nemotron_native/kernels/b32/codec.cu create mode 100644 deploy/nemotron_native/kernels/b32/native/codec/kernel_kf_codec.h create mode 100644 deploy/nemotron_native/kernels/b32/native/codec/legacy_codec.cuh create mode 100644 deploy/nemotron_native/kernels/b32/native/codec/native_stage1.cu create mode 100644 deploy/nemotron_native/kernels/b32/native/codec/native_stage1.h create mode 100644 deploy/nemotron_native/kernels/b32/native/codec/repack3.inc create mode 100644 deploy/nemotron_native/kernels/b32/native/codec/tail_kernel.inc create mode 100644 deploy/nemotron_native/kernels/b32/native/codec/wg2_kernel.inc create mode 100644 deploy/nemotron_native/kernels/b32/native/codec/wg3_kernel.inc create mode 100644 deploy/nemotron_native/kernels/b32/native/codec/wg4_step_widths.inc create mode 100644 deploy/nemotron_native/kernels/b32/native/codec/wg5_step_widths.inc create mode 100644 deploy/nemotron_native/kernels/b32/native/codec/wg_kernel.inc create mode 100644 deploy/nemotron_native/kernels/b32/native/codec/wgmma_macro.inc create mode 100644 deploy/nemotron_native/kernels/b32/native/prefill/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b32/native/prefill/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/b32/native/prefill/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b32/native/prefill/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b32/native/prefill/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b32/native/stage0/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b32/native/stage0/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b32/native/stage0/native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b32/native/stage0/predictor_mk2.cu create mode 100644 deploy/nemotron_native/kernels/b32/native/stage0/sparse_moe_impl.cuh create mode 100644 deploy/nemotron_native/kernels/b32/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b32/stage0.cu create mode 100644 deploy/nemotron_native/kernels/b4096/codec.cu create mode 100644 deploy/nemotron_native/kernels/b4096/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b4096/regions/prefill/native/prefill/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b4096/regions/prefill/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage0/native/stage0/predictor_mk2.cu create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage0/stage0.cu create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/codec.cu create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/kernel_kf_codec.h create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/legacy_codec.cuh create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/native_stage1.cu create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/native_stage1.h create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/repack3.inc create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/tail_kernel.inc create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg2_kernel.inc create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg3_kernel.inc create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg4_step_widths.inc create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg5_step_widths.inc create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wg_kernel.inc create mode 100644 deploy/nemotron_native/kernels/b4096/regions/stage1/native/codec/wgmma_macro.inc create mode 100644 deploy/nemotron_native/kernels/b4096/stage0.cu create mode 100644 deploy/nemotron_native/kernels/b512/codec.cu create mode 100644 deploy/nemotron_native/kernels/b512/native/codec/kernel_kf_codec.h create mode 100644 deploy/nemotron_native/kernels/b512/native/codec/legacy_codec.cuh create mode 100644 deploy/nemotron_native/kernels/b512/native/codec/native_stage1.cu create mode 100644 deploy/nemotron_native/kernels/b512/native/codec/native_stage1.h create mode 100644 deploy/nemotron_native/kernels/b512/native/codec/repack3.inc create mode 100644 deploy/nemotron_native/kernels/b512/native/codec/tail_kernel.inc create mode 100644 deploy/nemotron_native/kernels/b512/native/codec/wg2_kernel.inc create mode 100644 deploy/nemotron_native/kernels/b512/native/codec/wg3_kernel.inc create mode 100644 deploy/nemotron_native/kernels/b512/native/codec/wg4_step_widths.inc create mode 100644 deploy/nemotron_native/kernels/b512/native/codec/wg5_step_widths.inc create mode 100644 deploy/nemotron_native/kernels/b512/native/codec/wg_kernel.inc create mode 100644 deploy/nemotron_native/kernels/b512/native/codec/wgmma_macro.inc create mode 100644 deploy/nemotron_native/kernels/b512/native/prefill/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b512/native/prefill/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/b512/native/prefill/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b512/native/prefill/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b512/native/prefill/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b512/native/stage0/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b512/native/stage0/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/b512/native/stage0/native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/b512/native/stage0/predictor_mk2.cu create mode 100644 deploy/nemotron_native/kernels/b512/prefill.cu create mode 100644 deploy/nemotron_native/kernels/b512/stage0.cu create mode 100644 deploy/nemotron_native/kernels/codec/base/kernel_kf_codec.h create mode 100644 deploy/nemotron_native/kernels/codec/base/legacy_codec.cuh create mode 100644 deploy/nemotron_native/kernels/codec/base/native_stage1.cu create mode 100644 deploy/nemotron_native/kernels/codec/base/native_stage1.h create mode 100644 deploy/nemotron_native/kernels/codec/base/repack3.inc create mode 100644 deploy/nemotron_native/kernels/codec/base/tail_kernel.inc create mode 100644 deploy/nemotron_native/kernels/codec/base/wg2_kernel.inc create mode 100644 deploy/nemotron_native/kernels/codec/base/wg3_kernel.inc create mode 100644 deploy/nemotron_native/kernels/codec/base/wg4_step_widths.inc create mode 100644 deploy/nemotron_native/kernels/codec/base/wg5_step_widths.inc create mode 100644 deploy/nemotron_native/kernels/codec/base/wg_kernel.inc create mode 100644 deploy/nemotron_native/kernels/codec/base/wgmma_macro.inc create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/codec.cu create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/codec.cu create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/kernel_kf_codec.h create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/legacy_codec.cuh create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/native_stage1.cu create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/native_stage1.h create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/repack3.inc create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/tail_kernel.inc create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg2_kernel.inc create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg3_kernel.inc create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg4_step_widths.inc create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg5_step_widths.inc create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wg_kernel.inc create mode 100644 deploy/nemotron_native/kernels/codec/first_packet/regions/stage1/native/codec/wgmma_macro.inc create mode 100644 deploy/nemotron_native/kernels/codec/small/kernel_kf_codec.h create mode 100644 deploy/nemotron_native/kernels/codec/small/legacy_codec.cuh create mode 100644 deploy/nemotron_native/kernels/codec/small/native_stage1.cu create mode 100644 deploy/nemotron_native/kernels/codec/small/native_stage1.h create mode 100644 deploy/nemotron_native/kernels/codec/small/repack3.inc create mode 100644 deploy/nemotron_native/kernels/codec/small/tail_kernel.inc create mode 100644 deploy/nemotron_native/kernels/codec/small/wg2_kernel.inc create mode 100644 deploy/nemotron_native/kernels/codec/small/wg3_kernel.inc create mode 100644 deploy/nemotron_native/kernels/codec/small/wg4_step_widths.inc create mode 100644 deploy/nemotron_native/kernels/codec/small/wg5_step_widths.inc create mode 100644 deploy/nemotron_native/kernels/codec/small/wg_kernel.inc create mode 100644 deploy/nemotron_native/kernels/codec/small/wgmma_macro.inc create mode 100644 deploy/nemotron_native/kernels/contraction/bf16/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/contraction/bf16/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/contraction/bf16/native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/contraction/bf16/predictor_mk2.cu create mode 100644 deploy/nemotron_native/kernels/contraction/bf16/sparse_moe_impl.cuh create mode 100644 deploy/nemotron_native/kernels/contraction/lane_transfer.cu create mode 100644 deploy/nemotron_native/kernels/contraction/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/inflight/progress.cuh create mode 100644 deploy/nemotron_native/kernels/inflight/runtime.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b1/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/prefill/b1/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b1/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b1/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b1/prefill.cu create mode 100644 deploy/nemotron_native/kernels/prefill/b16/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/prefill/b16/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b16/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b16/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b16/prefill.cu create mode 100644 deploy/nemotron_native/kernels/prefill/b2/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/prefill/b2/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b2/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b2/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b2/prefill.cu create mode 100644 deploy/nemotron_native/kernels/prefill/b4/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/prefill/b4/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b4/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b4/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b4/prefill.cu create mode 100644 deploy/nemotron_native/kernels/prefill/b64/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/prefill/b64/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b64/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b64/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b64/prefill.cu create mode 100644 deploy/nemotron_native/kernels/prefill/b8/dedup_native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/prefill/b8/dedup_predictor_mk2.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b8/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b8/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/prefill/b8/prefill.cu create mode 100644 deploy/nemotron_native/kernels/recipes.json create mode 100644 deploy/nemotron_native/kernels/stage0/b1/bb_router_selection.cuh create mode 100644 deploy/nemotron_native/kernels/stage0/b1/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/stage0/b1/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/stage0/b1/native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/stage0/b1/predictor_mk2.cu create mode 100644 deploy/nemotron_native/kernels/stage0/b16/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/stage0/b16/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/stage0/b16/native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/stage0/b16/predictor_mk2.cu create mode 100644 deploy/nemotron_native/kernels/stage0/b16/sparse_moe_impl.cuh create mode 100644 deploy/nemotron_native/kernels/stage0/b64/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/stage0/b64/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/stage0/b64/native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/stage0/b64/predictor_mk2.cu create mode 100644 deploy/nemotron_native/kernels/stage0/small/mk2_abi.cuh create mode 100644 deploy/nemotron_native/kernels/stage0/small/native_prefill_abi.cuh create mode 100644 deploy/nemotron_native/kernels/stage0/small/native_prefill_impl.inc create mode 100644 deploy/nemotron_native/kernels/stage0/small/predictor_mk2.cu create mode 100644 deploy/nemotron_native/kernels/stage0/small/sparse_moe_impl.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/LICENSE create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/base_ops.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/base_types.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/common.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/multimem.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/common/util.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/kittens.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/group.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/memory.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_global_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_global_to_shared.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/complex_shared_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/global_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/global_to_shared.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/parallel_global_to_global.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/shared_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tensor_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tile.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/tile/tma_cluster.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/global_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/global_to_shared.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/shared_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/tma_cluster.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/memory/vec/vec.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x112.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x128.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x144.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x16.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x160.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x176.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x192.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x208.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x224.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x240.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x256.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x32.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x48.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x64.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x80.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/64x96.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/base/base.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/mma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/tcgen05.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/warp.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/mma/warpgroup.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/complex_conversions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/complex_maps.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/conversions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/maps.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/reductions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/tile/tile.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/conversions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/maps.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/reductions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/register/vec/vec.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/shared.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/conversions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/maps.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/reductions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/tile/tile.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/conversions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/maps.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/reductions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/shared/vec/vec.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/sync.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/tma_cluster.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/group/util/util.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/ops.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/memory.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/shared_to_tensor.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/tile.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/tile/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/vec/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/memory/vec/vec.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/mma/mma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/mma/tcgen05.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/thread.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/util/sync.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/util/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/ops/thread/util/util.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/broker.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/club.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/parallel_tensor.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/profiler.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/pyutils.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/pyutils/torchutils.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/cgl.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/gl.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/global.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/global/util.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/crt.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/crv.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rt.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rt_base.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rt_layout.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rv.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/register/rv_layout.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/cst.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/csv.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/descriptor.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/shared.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/st.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/shared/sv.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/ipc.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/pgl.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/system.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/system/vmm.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/tensor/tensor.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/tensor/tt.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens-v4-prefill/include/types/types.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/LICENSE create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/base_ops.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/base_types.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/common.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/multimem.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/common/util.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/kittens.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/group.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/memory.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/complex_global_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/complex_global_to_shared.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/complex_shared_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/global_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/global_to_shared.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/parallel_global_to_global.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/shared_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tensor_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tile.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/tile/tma_cluster.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/global_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/global_to_shared.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/shared_to_register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/tma_cluster.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/memory/vec/vec.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x112.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x128.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x144.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x16.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x160.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x176.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x192.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x208.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x224.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x240.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x256.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x32.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x48.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x64.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x80.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/64x96.impl create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/base/base.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/mma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/tcgen05.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/warp.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/mma/warpgroup.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/complex_conversions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/complex_maps.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/conversions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/maps.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/reductions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/tile/tile.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/conversions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/maps.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/reductions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/register/vec/vec.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/shared.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/conversions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/maps.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/reductions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/tile/tile.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/conversions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/maps.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/reductions.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/shared/vec/vec.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/sync.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/tma_cluster.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/group/util/util.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/ops.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/memory.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/tile/shared_to_tensor.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/tile/tile.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/tile/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/vec/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/memory/vec/vec.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/mma/mma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/mma/tcgen05.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/thread.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/util/sync.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/util/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/ops/thread/util/util.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/broker.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/club.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/parallel_tensor.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/profiler.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/pyutils.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/pyutils/torchutils.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/cgl.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/gl.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/global.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/tma.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/global/util.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/crt.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/crv.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/register.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rt.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rt_base.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rt_layout.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rv.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/register/rv_layout.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/cst.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/csv.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/descriptor.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/shared.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/st.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/shared/sv.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/ipc.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/pgl.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/system.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/system/vmm.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/tensor/tensor.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/tensor/tt.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/include/types/types.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/common/common.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/common/templates.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/common/util.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/interpreter/interpreter.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/interpreter/templates.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcf/lcf.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcf/templates.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcsc/lcsc.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcsf/lcsf.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/lcsf/templates.cuh create mode 100644 deploy/nemotron_native/kernels/third_party/thunderkittens/prototype/prototype.cuh create mode 100644 deploy/nemotron_native/model.lock.json create mode 100644 deploy/nemotron_native/profiles/playback-safe.json create mode 100644 deploy/nemotron_native/programs.json create mode 100644 deploy/nemotron_native/serve.py create mode 100644 docker/Dockerfile.nemotron create mode 100644 docs/nemotron-integration-results.md create mode 100644 docs/nemotron-megakernel.md create mode 100644 native/tts-normalizer/.bazelrc create mode 100644 native/tts-normalizer/BUILD.bazel create mode 100644 native/tts-normalizer/LICENSE.legacy create mode 100644 native/tts-normalizer/LICENSE.riva create mode 100644 native/tts-normalizer/README.md create mode 100644 native/tts-normalizer/THIRD_PARTY_NOTICES.md create mode 100644 native/tts-normalizer/WORKSPACE.bazel create mode 100644 native/tts-normalizer/include/processor_abi.h create mode 100644 native/tts-normalizer/src/normalize.cc create mode 100644 native/tts-normalizer/src/normalize.h create mode 100644 native/tts-normalizer/src/speech_hints.cc create mode 100644 native/tts-normalizer/src/speech_hints.h create mode 100644 native/tts-normalizer/src/text_preprocessor.cc create mode 100644 native/tts-normalizer/src/text_preprocessor.h create mode 100644 native/tts-normalizer/src/text_preprocessor_c.cc create mode 100644 native/tts-normalizer/testdata/BUILD.bazel create mode 100644 native/tts-normalizer/testdata/en-US/abbr.txt create mode 100644 native/tts-normalizer/testdata/en-US/post_process.far create mode 100644 native/tts-normalizer/testdata/en-US/postprocessor.ascii_proto create mode 100644 native/tts-normalizer/testdata/en-US/sentence_boundary_exceptions.txt create mode 100644 native/tts-normalizer/testdata/en-US/sparrowhawk_configuration.ascii_proto create mode 100644 native/tts-normalizer/testdata/en-US/tokenize_and_classify.far create mode 100644 native/tts-normalizer/testdata/en-US/tokenizer.ascii_proto create mode 100644 native/tts-normalizer/testdata/en-US/verbalize.far create mode 100644 native/tts-normalizer/testdata/en-US/verbalizer.ascii_proto create mode 100644 native/tts-normalizer/tests/normalize_test.cc create mode 100644 native/tts-normalizer/tests/speech_hints_test.cc create mode 100644 native/tts-normalizer/tests/string_matcher_test.cc create mode 100644 native/tts-normalizer/tests/text_preprocessor_c_test.cc create mode 100644 native/tts-normalizer/tests/text_preprocessor_test.cc create mode 100644 native/tts-normalizer/third_party/BUILD.bazel create mode 100644 native/tts-normalizer/third_party/BUILD.openfst create mode 100644 native/tts-normalizer/third_party/BUILD.pugixml create mode 100644 native/tts-normalizer/third_party/BUILD.re2 create mode 100644 native/tts-normalizer/third_party/BUILD.sparrowhawk_0_1 create mode 100644 native/tts-normalizer/third_party/BUILD.thrax_1_2_7 create mode 100644 native/tts-normalizer/third_party/sparrowhawk_normalizer.cc create mode 100644 native/tts-normalizer/third_party/sparrowhawk_rule_system.cc create mode 100644 native/tts-normalizer/third_party/sparrowhawk_string_matcher.h create mode 100644 python/nemotron_cpp_bridge/__init__.py create mode 100644 python/nemotron_cpp_bridge/__main__.py create mode 100644 python/nemotron_cpp_bridge/normalization.py create mode 100644 python/nemotron_cpp_bridge/protocol.py create mode 100644 python/nemotron_cpp_bridge/server.py create mode 100644 python/nemotron_cpp_bridge/slabs.py create mode 100644 python/nemotron_native_backend/__init__.py create mode 100644 python/nemotron_native_backend/backend.py create mode 100644 python/nemotron_native_backend/batcher.py create mode 100644 python/nemotron_native_backend/bundle.py create mode 100644 python/nemotron_native_backend/execution/__init__.py create mode 100644 python/nemotron_native_backend/execution/chunk_boundary.py create mode 100644 python/nemotron_native_backend/execution/chunk_submission.py create mode 100644 python/nemotron_native_backend/execution/composition.py create mode 100644 python/nemotron_native_backend/execution/controller_rules.py create mode 100644 python/nemotron_native_backend/execution/engine.py create mode 100644 python/nemotron_native_backend/execution/helper_graphs.py create mode 100644 python/nemotron_native_backend/execution/inflight_engine.py create mode 100644 python/nemotron_native_backend/execution/inflight_policy.py create mode 100644 python/nemotron_native_backend/execution/inflight_reset.py create mode 100644 python/nemotron_native_backend/execution/lane_control.py create mode 100644 python/nemotron_native_backend/execution/lane_windows.py create mode 100644 python/nemotron_native_backend/execution/narrow_prefix_cache.py create mode 100644 python/nemotron_native_backend/execution/narrow_state_binding.py create mode 100644 python/nemotron_native_backend/execution/native_assets.py create mode 100644 python/nemotron_native_backend/execution/native_inputs.py create mode 100644 python/nemotron_native_backend/execution/native_stage1.py create mode 100644 python/nemotron_native_backend/execution/pcm_transfer.py create mode 100644 python/nemotron_native_backend/execution/prefix_cache.py create mode 100644 python/nemotron_native_backend/execution/prefix_snapshot.py create mode 100644 python/nemotron_native_backend/execution/program.py create mode 100644 python/nemotron_native_backend/execution/readiness.py create mode 100644 python/nemotron_native_backend/execution/request_noise.py create mode 100644 python/nemotron_native_backend/execution/research_program.py create mode 100644 python/nemotron_native_backend/execution/retained_narrow.py create mode 100644 python/nemotron_native_backend/execution/reusable_inputs.py create mode 100644 python/nemotron_native_backend/execution/reusable_runner.py create mode 100644 python/nemotron_native_backend/execution/small_width.py create mode 100644 python/nemotron_native_backend/execution/source_build.py create mode 100644 python/nemotron_native_backend/execution/split_native_runtime.py create mode 100644 python/nemotron_native_backend/execution/startup_binding.py create mode 100644 python/nemotron_native_backend/execution/state_binding.py create mode 100644 python/nemotron_native_backend/execution/token_staging.py create mode 100644 python/nemotron_native_backend/execution/wide_noise.py create mode 100644 python/nemotron_native_backend/hardware.py create mode 100644 python/nemotron_native_backend/host/__init__.py create mode 100644 python/nemotron_native_backend/host/cpu_pools.py create mode 100644 python/nemotron_native_backend/host/diagnostic.py create mode 100644 python/nemotron_native_backend/host/exact_admission.py create mode 100644 python/nemotron_native_backend/host/first_fence.py create mode 100644 python/nemotron_native_backend/host/ready_writer.py create mode 100644 python/nemotron_native_backend/host/serialization.py create mode 100644 python/nemotron_native_backend/host/slab_audit.py create mode 100644 python/nemotron_native_backend/host/slab_cancel_v2.py create mode 100644 python/nemotron_native_backend/host/slab_pool.py create mode 100644 python/nemotron_native_backend/host/transport.py create mode 100644 python/nemotron_native_backend/inflight_batcher.py create mode 100644 python/nemotron_native_backend/initial_publication.py create mode 100644 python/nemotron_native_backend/lifecycle.py create mode 100644 python/nemotron_native_backend/manifest.py create mode 100644 python/nemotron_native_backend/models.py create mode 100644 python/nemotron_native_backend/protocol.py create mode 100644 python/nemotron_native_backend/r31.py create mode 100644 python/nemotron_native_backend/receipts.py create mode 100644 python/nemotron_native_backend/server.py create mode 100644 python/nemotron_native_backend/streaming_batcher.py create mode 100644 python/nemotron_native_backend/streaming_contract.py create mode 100644 python/nemotron_native_backend/streaming_models.py create mode 100644 python/nemotron_native_backend/streaming_server.py create mode 100644 python/nemotron_native_backend/worker.py create mode 100644 scripts/nemotron_benchmark.py create mode 100644 scripts/nemotron_native/build_kernels.py create mode 100644 scripts/nemotron_native/build_package.py create mode 100644 scripts/nemotron_native/launch.py create mode 100644 scripts/nemotron_native/requirements.txt create mode 100644 scripts/nemotron_native/runtime_dependencies.py create mode 100644 scripts/nemotron_native/tests/test_bridge.py create mode 100644 scripts/nemotron_native/tests/test_launcher.py create mode 100644 src/tts/nemotron/native_runtime.cpp create mode 100644 src/tts/nemotron/native_runtime.h create mode 100644 src/tts/nemotron/pcm_conversion.h create mode 100644 tests/cpp/common/test_nemotron_pcm.cpp create mode 100644 tests/nemotron_http_test.py create mode 100644 tests/nemotron_mock_worker.py create mode 100644 tests/nemotron_normalization_test.py 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..47cdf64 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,44 @@ 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). +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..aa974d6 --- /dev/null +++ b/docker/Dockerfile.nemotron @@ -0,0 +1,72 @@ +# 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 + +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 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 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 +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..f0ae03e --- /dev/null +++ b/docs/nemotron-integration-results.md @@ -0,0 +1,227 @@ +# 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. B512 +startup latency remains higher than the legacy frontend in the same-host +control. These samples do not establish production readiness, acoustic +quality, or a performance improvement over the legacy deployment. + +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. diff --git a/docs/nemotron-megakernel.md b/docs/nemotron-megakernel.md new file mode 100644 index 0000000..dc65e32 --- /dev/null +++ b/docs/nemotron-megakernel.md @@ -0,0 +1,328 @@ +# 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 bridge's tokenization and normalization. +The C++ frontend owns public HTTP handling and multiplexes requests over one +Unix socket. This integration is not a pure C++/CUDA rewrite of that controller. +There is no TensorRT-LLM execution dependency in the retained native path. + +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 + +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. + +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. + +To reproduce the optional float32-forwarding benchmark candidate, add +`-e NEMOTRON_BRIDGE_FORWARD_F32=1` before the image name in the `docker run` +command above. This moves float32-to-PCM16 conversion to the C++ frontend; +public HTTP audio remains PCM16. Leave the variable unset for the default +PCM16 bridge path, and record the selected value with benchmark receipts. +The [results report](nemotron-integration-results.md) distinguishes this +candidate from the default and records its qualification limits. + +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 +``` + +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/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 0000000000000000000000000000000000000000..45bc69d1ecfd1dcd8f0e6c09e27604f6135c23c5 GIT binary patch literal 327407 zcmeFa1-PB%xwT7hcPGGyy9Y~<5C|by3@1RawiJpbxFl_9sS9#IPrfeIuWsz9=l3M3n=K(dJnBtujn*;ECR z%~T-STm_OXR3O<>1(KmEkPK6SWGfXQ9KWp%PPQ>k^kG{QAFcw)b}B%4oQ^Q@?Nxv< z-@(L3ssLfWqlu4F0m6JI6CbStg!#@UzKaSF=DV8sZYn^S?{4B_RDduaYvSWnfG{6# z;uBPWFyF((_f!GGd@mEm9H;`xK`M|;QsFVq zR$@Re&&mkBPrm1qk!^nfMVZK$ySZ z#E(<~!u%)`|9}b*=0}_OF)BcqA8X>rsQ_Vqyopa%0m6KWiJzbXg!zdkev%3h<|mu@ zDJnpipK9W#sQ_VrINjjn4AVp(&NT6g!yX*AH;`QrcGVAt!kJ%zxM|XpTa+6;vZH4!u%`~|A-0@<{vfjvsHjF|Cot? zTm=a8Pnh^8Re&)6l!<>@1qk!cnD}Q^fH0p$$6|O})i8HvGuyHFoc_!}`160>#LrOy z!u$&+{zVla%)eyfUseIa{3|B@RTUu2zh>fJR{_HO8z%lu6(G#NW#Z?m0Ac=Z6Q8aE zg!y?U{v8z{%)e{m=c@o={yh`_z6ucL7nt}DRDdx5p^5)U1qkyWoA^&ufH425iT_Lm z2=kwt_%Bp|F#n~AU#J3v`L9g;*D64m|Hi~GQUSvJViW(Z3J~VMGx6W60Ac+5-uc|9-U7Z@$=_Vv8Tyq*{FdS1lqc@eMYMZBID@p@jw z>v<9Xh0)T!pL$+oUeAkoJul+*yolHHB3{pncs(!T^}LAJ^CDi)i+DXR;`O|U*YhG? z&x?3HFXEgRySDcCf|-3UsP~2Gx=#S%e$w9;de_W39y7zfFq}WEGR@FQy+*+6H3D9* z5%7AAfY)mTyj~;VXB#@H*9dsMM!@Se0$#5X@Oq7a*J}j4UL)Z38Ue5S0$#5X@Oq7a z*J}j4UL)Z38Ue4@2zb3l!0R;vUat}GdX0eBYXrPrBjEKK0k78xc)dozryDw{djVdr z5%7AAfY)mTyj~;V^%?=cz|cv(M!@Se0$#5X@Oq7a*J}j4UL)Z38Ue4@2zb3l!0R;v zUiSjLUL)Z38Ue4@2zb3l!0R;vUat}GdX0eBYXrPrBjEKK0k78xc)doz>oo#S_@2lh z(?lO;HSyV0AemhSk~vf$nNtN4&NYzCtpdqBDv->p0?B+Tkj$?F$pR{nET{s>LMo6f ztOCg*Dv&Iy0?A@3kSwkO$r37%EU5y?QYw%vtpdq1Dv&Iz0?BeJkPKFVWO)@xR#1Ut zMHNW86s`Om;(336C4CN(l~o{#_utR&9g$76Kr%!Hl1)`0*-QnJ%~c@TLIsj7RUjFv z0?9BHNVZY|!q0bWgOhDc6Mfj$#D}XuvYiSL9;YKre0voj%y%&Hkt#r#?`YzqRDdww z$;3yi0AaqfiSMEUg!!%}zMBdV=DVBt7!@GQ$C~&!6(G#VoA?A3Ak6nL@jX?5FyG6> z_f`SId><3vR|N?3{Y-p+6-XwkKyrWzBnPTMa*zrnlT>(&_utR&9q*%cEP!zQk1_FM zRe&%*&cu&b0m6K;iBC}h!u$jiKT!n;^OH>cWECLHPciXRRe&%*&EV`0ryHD{VVdZ} znI=9}1(HitfN-9_)Wk1S0mA%p6Td2nG4WSbfG~f};Dh*}{I<9Buf25~ z7V)uQS=8iHi>W}exC$gos6eu$3M5OZK(e$7B+IBkvaAXu%c(#zSOt>hRUlbG1(Fq2 zAn9DV*5-C`KF#fLAbj1G3{F-yO^Ux8*W&y+*EIZobBnbMPS!R}^kE$nUsnZ^^;95P zUj>p4R3O<<1(JW50VSb#6AFl$0 z`D7EHq5_2Z2_}A`3J~TenfS>nK$xFm;-{(rVSbvy*&j|fI61>K(T6ile5wj0m#6?) zMqFy}xx{5Aez^(|=3IvZ!LKy&t5kq6zuLsP4hOQFxYppyit7x{ufN{le7`psd}(o` z!MXiSCVsOD5azd-_^m2HnBQjNx2phIO59;^Zhxo2mlSsyocY}*evb+e=J%TTeJVhh z-*4g%r~qO9puv|A4;d|gbuGTQKKF=WzLl0Ac=4UW+fPulKCc@}IsI zUqoNyCBxTSSiEfTS;Z?R{;CQP=C2ui5Fhkj`KE0hxb8x^?yBe7PCmkDnWYo;T#N5u zw6yt1tw1_a&$aj{qovL3xwg~iMjI_{UeC3iK3C7Rcs6 zv;yhGc%!BH1g$_iv4_#pd{3=FImG(TJ`kWReEXledltw1{QKBJ}i5n6$C;{8TT^CPta>BLb+OY;wC1=5Ln zZo})j4L{asX}3R4E09hcZ?rU@tQANnrWh^FPtXda6P(w8bmAnVr8(y{Ae}hHXlZ_` zRv=yb*n8iMI>`9b?Xk$?F_}#T680iU=1_rTP8CSzQh{V{6-ee$fn;74Naj<4WPTM$ z7Epm?K@~_AQh|iM36e!rAYpHWWHA*;*vlbVLIo1eMUX6|0?E=UkSwDD$+9YtaIS`A zunHv0t3a}X3M4D4K+?J7JtaJ!W3-Z?la)=A;_t@uT7Hi;4W0YRS_UU;n0Jn0)%-zcjKJ5f#CJr zjdR`xg1^tik5B=^{QV|=qzVw`N16BsRDdu)+Qg4h0mA%P6F*J`2=n7je6k7<=2J}k z1Qj66Pc-q9RDdu)*~Cv#0mA%L6F*G_2#@d64NlH5P4wYR6Q8OA$t5a4I8R?{;+Lra zVSc%ZU!eko`IRPql?o8%SDW}XDnOWDYvR|b0AYT;iQk|Cg!zpoev=9i<~N)8Eh<2m z-)iEwsQ_VqyNTbS0)+XUCVrO+5axHA_&q8>nBQyS_o)D3e!q!7paO*XgC_ov3XnO) z!v>#SJYwRHssLgBn2A5G0)+V!CjO)f5av&r_|qytm_K9U&#C}n{+x+FuL5Ka@q)p* z{fj34k_wR7#LEVsRlH*2uc`oH{+huD@j=f^d+RK-OlscyuwYx>eAc_gn?IjjMZ@h3 zw_90^F!AkGfH2>|#7C+CVZNh@k5U1`d?ynhtpbGk&L+N#3J~VIn)q%iK$!1t;$u{R zFdu8;<5Yk!A8+ClRDdww!^G=3mO1BFAb35;;+$WB;Po7fbAAPaA8g`>r~qMpsEHq@ z0)+YDCjK53Ak5!u;_p)d!u$vmf4>S4=0}?NQ7S-~f560#Rsq8N7!yBM1qk!wO#FBi zAj~J5_!Jc&%iK44Ip-Q68;N?p*-*~82FM1Y zo^SYb*7M7Ha?UM4I4<@4v5x#9qoqCfdRQwEj_o4`UsF74@HIr;-&dEvXtcC{OXwx7 zKvr!Z|2hBd=w|x8zhd|{W*1*I@vo@>Vg7X!|Aq<>=HE2&Z>a!bey)jsTLlR7=|)TQ z^RxnC{v8wlt_l$5=bQNVRDdx5zKLI;0)+VwO#Fu`K$!o?#DA;;g!xZQ{HH2FnE%Yg zf35t2A@ zy#TL!0bch4yzT{f-3#!#7vOa-!0TSHg`9l=$mXK%1>C;w1$ciiXa}L)7iXruU`ze6 z`+LE_bpelsf2F;EKTq8Y@PBkK7^+{}zb@!s7vwnt_oe=ILGRi?hoYT3X5{OFcXBV_ zYxVbnfxTdB9gi9LIRfYT{&R%h^8_7=_E?>npCkOey1=JoXgK388a!0YP;czwM9udf&2_4NY0zFvUW*9-9a z909Md7vOa-!0YP;czwM9udf&2_4NY0zFvUW*9-9adI4TvFTm^T1$f;H@NXG9sjnB{ z^*I7Q-Ovfw1i3E2>vaKMUoXJx>jgNk6+lv7FW_?*n9S?z1$?f)UVzuv3-J1S0bXA( z!0TRs*Vha1`g#FgpCjP)^#Z)UUVzu@0=&LnFj!t+FTm^T1>C;AUVzuv3-J1S0bXA( z!0YP;_#X|O)aM9zeU5jJJT z-s-hO=JmBhy#Lx^@Abpp(eA%?I3vAw$j_v|7v!G#`oDL`XZ!CX^xjLz5$eB>F!29! z;+Xt*ejkDF|DAljfOFkjy^p~Bk48)T@5bx<2+ZsI2>4$=N9g~XcpN*jh-sn^i<cWECLHPciXRRe&%*&EV`0ryHD{VVdZ}nI=9}1(Hit zfXpQ>HSx<-fH1$@#IH~R!u(1Tze)uN^Q%q#8WkYSuQl=ORDdwQ-o$TE0mA%76Te9X z2=kjw{1z1;%x^XE+f;xszum;|PyxdHP7}XN1qk!IP5d4eAk6PI@%vPOFu&i#A5a0p z{6Q0cNCn88;$ef&E*>%QM^%6@f6T-mR{_HO2@`)(1qkz}O#EpTAk3dJ@n=1L&EO|3Ev+i zeE*Q}^MQn)CnP-HAz?p&WL_0W=2L-WeicX-P=RDY6-X9Rfn;G7NET6nWKk7J7E^&_ zaTQ3GP=RDg6-bs+fn;eFNS0B7WLXtRmQ#UbunHv0t3a}X3M4D4K+?HzExwyN4#`R? zkgTi%N&MZo_TE*$$(n}WkNe451}AHqCi<|BiLa{y$$Bb~tgiye1}czjr~=7GDv)ff z0?8&SkPK0QWK$JLHdBFQa}`LoP=RDi6-b7vKr&1PlC4yLaQwD5IN8QD(T8nKe7FiE z+o=HIaXP}pw^sqed{RDduaZ{ibFfH2>~#P?JI!hA0i-&+L;^LcWECLHPciXRRe&%*&EV`0ryHD{VVdZ}nI=9}1(HitfQ%NG8l3a~WhQ>P3J~U3 znD~_{K$u@;;#aExVSbIl*$1vQID5r)2IuzI8=UX=27|Nr+-Puaf0K#dtOA7jEhc`e z3J~VEnfUE0Kt_o>49@NEG&pW<$gJX36Mszw$e{Ld=J+>;#_7*m@7o*6>wSBy zyxzArkk|Y6`to{TUQb@{%j?SPeRUoALq^Nr{Wpfj=r;Afx`w>omsXe8`@-(>dS6&o zK8yZ4sfF4)a1G6M^+Ea^Uk?c9+e-}2`MT~~oS*ByHBnyoDQ;i)sr}@2AL06-9@~B7 z_1Nq!f7obg*C3B*1+u58fA&4(_0K**UjOXl<^0(%YU}I&?dRk5IUWZ<@Yzh9#~Bdj zJPv{2bDH>EDnOWX-U5QpW8(9w0AW6#iO;VBg!uv{zMu*a=KPHFbH*1oTH1UOtw5MB zYT}Ej0Aaqki7%l7g!z&tzLW|O=1ZISGAclrFKgn9pgRe&(x%)~cW z0m6I>6W>w=2=k#PK1>A&^Q{a%S!``^=G&O~wkklF4>wwxZ>JSVCq@`8&9~PIq!T+B zEzL)21=5Kfjh5!4v;yhGPDV@f(OQ9YVrQeJ`7T<4bYfSdrTK1JfplVbqow&6tw1_4 z)@W%yPAia3j5k`EPtXda6YRf0I0`~zBnbmC~ErTH;hfpp?n zqow(AT7h(eV+y1bJO+VuVv5nyoMQ~66YPONI>8kQ8AuQxc~?+pg$XMLl=x&2KhezOV?=C_#mttvp6-)7>ss{rACc89^a{hbEq zzIm6yncr>V_ox72ey@q&rvile{U-i^3J~TG8l1<@Lqu-49?^7QG@f? ze$2!lR{_HO2@`)(1qkz}O#EpTAk3dJID5sj250`9i9fFb8i| zFB|+|@rsGRsse=hYX+ah2VH9qX>0F1x}AHcb%fzlo!H)JX}*J2Ae|U#v^3vQE09i% zGFqDNq!mafMjI{7ch(A|6T28K&3Dxbq!YUtEzNh=3ZxTbjF#qOwF2qHIHRTcc&$J> zF~MkQzK2#Io!HZ8Y0h~RNGJZY=hB1qHGp)2pFfaJaQ^_(3GP2YI>G%5NGG^|0_g;g z7a*PB@dTt3M;I;5c{~E?1dmrBonXHJ(h2q#Ae~@;1JVihKOmhr)@W&doK_&6INoS! zK3OY}PE0XcnsXil(h1IEKsxcCJ(uxVx|lDf3^MdjiCIm2HWeVuXE*UVRDdv_)5Pae z0m6K46Q4%~2=jSOd_EN*%;z`p1yq1AU(m!CQUStzVH00O1qky+O?)vGAj}sx@g-D% zFkjNdmr?=3d}$M3Mg<7-WlelJ6(Gz9oA~l7K$x#!;w!2EVct;`u0zk!+9A#tw04G1 zeO|0&aDM&D=8ewR`;tDniphLc6(Icj)l7VK6(G#lF!<-hng;*0Sj*tter*$9M+FG; zbxnLd6(G#lH}MTrfH2?C#5Ym_!hB;B-$Vrn^C2d_sR|J0o0<6LDnOWTVd7h=0AW7V z#D}Q>VZN2YIeuFkocT5;zO4!n=EIGa=G$on!hK6n7zQ4iw`AszO15|)8KhVSvQUStzlHlu-&x%O`=kfkUJD($; z5eFNb+Z|%!hpGT!ewc|Lt^$PldrbVjDnOWX-UosoVd9+kfiOSP#E((|!u$g!ezXb@ z=Es=$u_{2AA7|pns{mm>*~F))0AYTDiJzzfg!xG(ezFP>=BJqWsVYF2pJs6OgVPPp z&-)A$KT`z=^Qk6&i3*U9i%Sju5pkJ`U#L$H`E4eCy9$uA#T^Fc_IDcmqv9@u zGr!x!?@gZCbXZ~i=XWz9zzZi{bk;yb7SVLsBtcT@qwe3XgrqymKbXcOOA z1qkz9Ong@rAk23&@!eH`Fdt*$V^x4KA7|p@Re&&`VB&kI0Ac=apG*7avVrHZ9Pi_e zmUjFnYX!o5iiyA5=Q5th&c$!0e9q9nDL!xTZ-{dY{&n$XgMUqM8z5g5-!k}D#JL9N z*YbPM(bm`hn`Yd%~6{h9xo8|lxoXj@h(VGKjgnKTH1dT%g+Z0_sx22|J*)1ZR!GTef{^wwGp=m!nITVUT>GzpWAKndVjrD zUhl8`-0Jsoi=6uskedZRYalm?`g6Te&d&$P4TA5#_h;x`YyL$J4DKz~}0H3_r`z$wy2RefX$}pREGP$5bHsxC$hnP=Vx= zDv;F2C!gc~4g(qZ|37|FpM-%-n>xFt-WT~){rTeTk&tjdg@Fv5L;l{M;qQIV^>N47 zXK#n3KJI#-8#sqd(XMgX{RqwOCUq5Z?oc=;}Z{bsK z{4>_Qg?Zg)@&Cv%7{m>gH{Qp%zi0Ml-{Ty;U-r5^flu+chUATpGxmQ->SMI`xoJ~} zdpJ!ikhw&CoZEO*S!F*djVee0=(`8c-;%|x)FBrHk;Q06Vf`Qj% zI3E4Iptm1zEXcp=UcjI2o!kr9Tl&`pc^#Ic+rKX8Ul$C#rpNtZM!hcJdwVDM0>0jV z=e>ZvtX>!3Z}l93`8SQ0mi}{u*S}u+PVNO9@3-?Df!p_=BlPZjgYk|CL`a;LlfIFTnfP1v&^$|Fy&3>xVfW(^VkpzjiqAcK|ph|D9hu zeB)lgpQr8x%x9F>4*5>r&T|C5_x}3`z4sD0cI2C;i9Yn7BfNXh5je(g=Q#qmpYfj~ z4E&vY@!0m5EV!^Re@wP6-YK$fn*C6NVZgg zWT*-x!&D&IN(BhVZ)<~-ZA=q=*w(~{t3a}y3J@NrBTRgI6(G!aF!7NpK$!1n;-gf6 zFyG0r~zO#w%q5_2Zt|q>l3J~VIoA?+NAk4>__&60H%*UJf1Qj66_b~B2Re&(x z%f$Cq0m6JA6W><_2=o0+e18>4CaOSkfC?lBsz7p(3M7+Mc#QY|?a|)xK3c~D2*>{z z6F*i32=n7i{CE{0%qJWCZZXB+%ug`!6IFmPKgq;TRsq8N6cayH1qk!g49@;=y1~gA zrinhBY2s5=Ah|>Z2S0tABfx>%d10bNo>ge@q1k^T$p62^ApBpEU8ORDdvl z+Qgqx0mA%Q6Ms$x2=nI+K8JY0;LKk%@t0J9%qCto_^jd;6Mt0&2=mtrK8O$H?-bX) zfS)6?4cq$qzw5^L^?Ab>KgYzspaO*X7ft+2DnOWj+2Gt?zG86ZUp4WssQ_XAbrb)F z3J~VsH1Th#0AYTviGN!K2=nPiOY`%z0%86g6aTIX5a#Eb`1e$RF#o=ZU!Ven`43F| zhblmr|H#CDtOA7jPfYx$DnOY3%*21L0)+W5O#GKBK$u@>;=fV>!u;1J{u>n_%r7$W zi&cOy|E-DtP6Y__-<$X!RDdx5qly1X1qky$oA_T;fH0rsZ~VR2xitBp;a2#EO#H(t zK$xFp;vZ1~!u+Eqezpn_<{vZhkE;M-{s|NRqzVw`pEB`Js{mpC8594k3J~U>Gx53? zFt2+7UiSjL?ge<=3-G!Z;B_y+>t2A@y#TL!0bch4yzT{f-3#!#7vOa-z^5B6?ReC^ zfO*{u@VXb^=Nm2U*Ru}*VP5wF`~st;&FfykyzT{f-3#!#7vOa-!0TRs*S!F*djVee z0=(`8c-;%|x)Ry2VqkF-F`nCP*g5L849g4PB^sftg z*9JKr{p*5(=LkI3W~A$acXBV_YxVbnfxX}n9gi7#FW@}ie~!?5o}fd~_U@VaIl|xD z3m(=ns?QM~;-LMF*9(}`y#TMT7vS~v0=&Lnfb&`bB=z+IK6kdsyuMz*=j!VPczwM9 zudf&2_4NY0zFvUW*9-9adI4UaBjEM*0=(`8czwM9udf&2_4NY0zFvUW*9-9adI4Tv zFTm^T1$ceE0Iz!iUSBW3&oy*XpCjPY4V|24n&?A)y@1cv>jJ#KUV!sj0VMVH0zP+v z$-KT^z~}1g1$ceE0I#nX;Pv$ayzT{feZ2s$uNUC;IRai^FTm^T1$ceE0I$~tyUOe9 z1$ceEfZNyC3-J1S0bXA(!0YP;czwM9udf&2^*I7wpCjOPFTneILGLK=aYo(?xG(kh zg5G|RpV*AF7x43`dja2L-3##kUcgc4rGH)U&i(f=-1pwjbpf~UUl;VQ4SaWeSFa2B zJ^Fh=??`jsjixM{&NJ5LN9OUx`6A7w|ec6d425=|CG_v(tqu6;O`xBK>F7O|8>^|+#mbT5!%DB zeds?&82CHL9FYF|2))01n0v^-^7{z<8Q;mhfW76d-bY|w-$%gf`v^GC5AqxV|Lf-n z1AixakmiugsshPuDv->s0?8aIkj$w93GaDAGPepO^Qb^FuL>mdsX#Kn3M31tK(e3; zBnzoPvakvyi>N@ds0t*DsX(&03M5OYK(eF?BulA4va|{$%cwxItO_K{sX#JV1(M}e zAXz~Lk`+}T>0I9IzmwfYvz3f**8ZEhl_Of|(Esmbdk19`EszXRfn-w^NH$Y}WOEfr zworj&OBG0lsz5SK1(K~)fbclm+TetrF(linKr&nflI>K0@Hib|;@hhLVZMWjk5mD| zd`A->r2>TcP9{ED1qkzAk5kCfZ)5C`0gq|n2#~>u_{2Ak2CS{DnOV|F!4Q9 zfH2?F#P?DG!hCNN-$w-q^LZAUKJqB-)G`Sr~qO9eiJ`Z1qkz_ zO#A~XK$stG;>V}}VScQMAEyF@`SB(`Sp^95DJFh`3J~Ten)pd7K$xFw;-{zpVScKK zpQZwY{o!ceK^y^r>a15i3*Up#HA*FnFwSA8dA)CQAE@{34dnH{&7ZU0m)DbXKL)~asrS`&hgMD;I{R?u&SJ&CHEonYePRz{MO*p#qW%k_V@b!s1*qJ z_dgl@TYS*hp1-YXm^+V$4;nsoq4*{@#Q2$$Ff7w|Vpbrxv&UhUwkz zt^fWdw?Eh9_VqaRJ~#0H&HA!F0i?GdO`AGvPha=^oV}?2Y`KqdJbHV<>wnMZ-Y}cS zE!*d^XAS)S**>q&1L1yI|EzqSKN~IWe*71$KwfAc2fp8WOb^to{(N!1-~8FWCIiB( z{(Rne+plX5gju~G=C(NwXK4=Pjk7m?e|$f)_ip;S|Av`g_X)!o|D=h(+1JWrxZZF1 zwSO@AwLD&dFt0yX{LeRC4OPhaHD-h;&@4)N+fY&_?|EAH>etq3Pn4fDhulqIg=_d2@RDdx5j){L)1qk!= zP5gT*K$zF_CSK2*_z#ShcKaV{1;YGCCjMg;Ak2Sa;y+aZ!u)3@{&N)|%zt6xzf=Lj zoNEUlc-_nKx>w@AF%(+vGcQO}Eby+*+6 zH3EK?(b9fB*9SnD*J}j4UL)Z38Ue4@2zb3lz&~lUw69;U5t!F&1iW4&;Pn~-ujfU) z?ge-~FXHvQh}ZKXUeAkoJul*QFTm@05wGV(yq*{FdS1lqc@eMYMZBID@p@jwryDKp zc-HeG^Lk#y>v<7B-)L#So^v7)=JmXY*YhG?_X5107x8*t#D8qGw69;!i_GhJ5wGV( zyq*{FdS1lqc@eMYMZBID@p@jw>v<8c=S94p7x8*t#Oq#w*YhG?&x?3HFaB@xdR}}@ zUeAl%zMdEHzjckwtiKlwJg??>^!I{+e`|u{(ccSt`+*Kc+lywTy?~!j-3##l$X?Lt zc4V5N^K1Ki!OXQ6^v?_2^LzLGndutwf7hS6e~mcs+6jB#ztT10KYCsmuH!(a8T$Vr z`san2>%1^sKXV}5m-^2a2VO_zc>H&IzW7e=1$?dkUNEp1{NFkrGxGCA&h!1(jt2fc zGVZ@K^J_!MRuMyw;z4qgE|ET8$yuPP^*Y_0g`kn&* z5u>I39oqVy0`vNw0$$%!!0USoczsU+|D@5w60L zIYvwS{px!P%w5}#eNO?e?w5}#eNO?euOs61bws?L7x2IJ+7YwK zALRZ~_X3Vje=p$3^d39?>w+c0R z2nxrd{~Tf9@8fbjzHGF#nAjo+97RJNbG6 z=eoCgPmy_jPZ6*0DdP1#Mf}a4BV44lm zE1CGpDnOX8V&bc+0AaqGiLb5#g!vjKzNQM0$zm;obNjVTd>s`a%-1#X^;Cc`U*E(x zPyxbxLlfUf1qkzvO?(p-Ak2rD_@*jAm~Up{o2vj}zJ-ZzsRD%gP!k`f0)+Wi2Iu%~ zZE)t>nE19TK$s6VTAFXC6-XyW7%k1W*9xQ)I~Xm^M`{Jqi5-oW=A*O%>BLS(OY_lM zfplVLqow&ST7h(8SEHr*Zd!qKVt1pZ`53K0Ix*I0X+BOXkWP#@TAEML3ZxTz7%k2B z)C!~%dl@aw_tpxe6Z;q~&G*#`q!arYEzS4W3WT5EL=!(i1qky4P5dAgAj~I;Jibq` z^Evj5`HHKye`E1rtw1_)h|$vgP^~~ZahTE4{BW&6I`JN(rTKfc0_nv2jF#p{Xa&-V z_ZuzEkJJjJ6Gs^>%|D)d>IuW%$GIs$?2=mQMd~+2b%(pP{EmeRpA8O*mRDdww%HSNo ztqsn68x!AF1qk!uMoaVUv;yhG2&1L>_F92-Vh5w8`ADrmI}s?$-%Tr!PV8>9G#{fCNGHY`EzQSi1=5M}MoaSvT7h(852K~| zo?3x)VlShm`QBQAbYdT)rTM;EfplU&qow)&T7mHMn`q((r~qMppot%(0)+V_!DF19 zASRjqwRP`!AFJa5q!Y&(EzLQGKsv!=5J)Gc7%j~?#y~o8qS4a)B&|R?akA0U{1mM~ zI&rGe()={7K=?jRH#pzt876+F3J~T~P5crSAV-QzP5d$yAj~f}ajvU@Fu&5oxvmDn z{Av^Dx*7=cYfYT%Y9P$7H*v13fiS<(#BWjo!u)0vzeNQI^IJ{)HWeVuZ#VHfRDdwQ z)5Pym0mA%l6Te3V2=jYQ{5};R%tD6Ce{H>pjsp<(mPJi`F%=-p7dP=GRDdvF(!`fi0m6J~6JJIJ z2=ir4d^r^$%muVCUUssLf$QIuRucg=bq{#^?WgxjxV;+%(pFki*QS5*PR zd^HnaT?Gj9HB5X>6(EzvS_bF#Yn%8wDnOX8YvSvv0AaqqiEp3+g!zUhzL5$L<{O*% zCMrOf4>9pgRe&(x%)~cW0m6I>6W>w=2=k#PK1>A&^Q{cd@!Q(q%(pS|ZB>9UA8xcX z-%cx#PK+>Gns2WaNGEnMTAGj43ZxS|8ZFI7X$8`Wos5>|qqPF*#Lh-b^Ifz8>BO!^ zOY_~d0_nu=MoaTCT7h(8tkKeZoK_&67;m&RpP&^;C-yK}n(wIk=qZLRejx}1E zAEy;aCpacRI>BQKNGGNkEzM8R3ZxS!8ZFIF(h8&#CmSuzPtgjb6Q>$2%}>(`gzw{Y zgY$i!Vd7`10AW7W#4k|+!g0OS#4l3;!u)a*zd{8F^D9mKDit8iuQu^(RDdwQ*2J$< z0mA%x6Td+P2=g0F{3aD3%x^aFTU3BBztzNVQvt&Kb`!rt1qkyyP5dquAk6PJ@q1K& zFu&Kt?^6N7{C*REKm`c%2TlAT6(BsW9yU0S+eb|NQ57J}A2adCRe&&m!o;6c0mA$# z6MtF-2=ix5{8<$s%%3yy=T(5P&%9u8ZvUc*zoY`>-^9xXKUlnC;;*UzVg8!IC-Fho z&V$+-?+K4BW*7Bz7H*T#c0!tEd2BzN(3@ zrUHcd>L$L13J~UN8ho-?%izq{Ht}^-fG}Uz#Me^+!hC%b-#`Tj^9@aWBNZUbH#YH2 zRDduaV&a>s0AaqFiEpj~g!vXGzNHEf=0i<77-6(D-(D+_PV8W`G#{xINGEnQTAGj23ZxS|87<96YX#DYosE{}yJ!W{iCvAB z=DTSH(uv)TmgZx$0_nt9qow&ctw1_4-e_q)K`W3>>|wMt-%~4)PV8m0G~Zh*kWTDl zv^3vWE09j?XS6imUn>xPeiKdn02Lt24>a+ERDdv_BzTOI6T~Fbe+|2LypPrK0Md!$ zjF#pcLm-{tF$km+Q;e489AhAzIMHZnev(!oojBQOX?}`UAe}hXXlZ_$Rv>&IryHE_ z^9&O|Qw0d~sV07j3Xmhkr6zuv3J~U(oA?zfK$u@?;#a8vVScrVU!wwq`L!l~oeB`< z*PHkaDnOXuXyP}i0AYT!iQl3Eg!!!|ewzvq=C_;p9V$SW-)Z7^sQ_Vqw~61Q0)+X! zCVrm^5a#!r_ya0Hm_KOZ52*k-LOg8n!^9&d{-_EN=8u{9<0?RyKVjlessLgBl!-sB z0)+W9CjP7n5a!RB`12}24i_&NoZG)>;xDNH`8V;h!4DR%nE0zIK$yQ~@JW2=Uu(Cc zz+SS5;S>0xCcc;o5ax@U_!25Wm@jGKOQ`^1zO;!iqXLBavL?Qq3J~UlO?-J3Ak0@V z@fB5oFz+bNziaK>IhQUlsOR|{w_nL{Go14<5PTIAUsVMN^VLjzbrm4Y*D(0}`u*3` z3S_ca%i!F8Z4+Nd1qkzXO?*8SAk5b{@eNdfFyGL`H&Ow@d}9;eLBJ63OY@OhfplU=qow&Mtw1`llhM+Av{oRU*x6`lzKd2Mo!HfAX}+6QAf4FVXlXu1 zE09i%HCmdF(+Z>$BL?}OY^<80_ns)MoaU3wF2qHenw04 z{j~z&=Qq*B4^RQZ{6G^wNCgP~yiQlXOg!wHdeya)) z=C_&n?J7X#((iYNRv_H|PJ_>>`CVFpFu&Wx?@K$yRi*V?n{_kUJj69`}bITL?g1<2vz z1%q?@7Y#m0U+*QY&?0|X1<1kT6%&6|1qk!k3_ghuy4D`t*1x`%dgHaty1E?@9((JX z_)rxf%!iryHYz}vZ)@TcRe&&m<72-*wpY}jy*_qUkk`k~^74A$9xSi-?d9b4zP+rx z-nW;L*ZcO;@_JuhO3wWl$ddAUUtL1}kUkG&arwhqfh;Dk_tiz^^}f1@ygqMRSYGc7 z3(4z!VL>@R%Y|8#u6^vC7uVBl4LOi?<&(5R&YShvuBG`kngLlyeyvs@Ys>4g=l1p3 zuPLv`ehqm&)~n0wv06?3u)Z#kRps^AtRk<+Vr6+f7Awi?vFPOWSWLqNIa7Sl#6P40 zg!zX}{45n9%s*n{A5{Uu{A?5dmHu2A>0Ac=F6X#eA zN0e!X&TVJ5WAQotnSt=<|GbHxqXLBa7fk$%DnOWj$;7{`0)+WjO#G`VK$w5c#J{cr zg!wm2{F^F3n19Q}&s71!{M#lzT?Gj9^Gy6ZDnOWj*Tm0P0mA%yCjNaDAj~f?@gJxF zVg5rC|B(t1=07&^pQr#~{!k`+SjY+ zMdtOqh}XRUujfU)o)__YUc~Ep5wGV(yq*{FdS1kTVYIaGx1JZ7*YhG?&x?3HFXHvQ zh}ZKXUeAkoJul+*yolHHB3{pncs(y}sB=v{FK!_Jvu*=~+t>3V{!=)`8Tt9*JGmF|wfcL(z+SMa_Qo0c`6B1}{%c1A@2PPAota-d`g?o95ZzzuYe$>N zd5+MrDANs{uXUb@_sw5}# zeNO?e?E#!0USoczwQz*Y_0gdS1ZmdkT1cPXVv*Dd6=z1-!ne zfYw5}#Jul$(Jq7$+qop06`h1c3bfcxs>t4XT?ge;#9TBhZDd68T zTH4pE?OyuiG^r-0Y@6!7|d5wGtl;PpKP zyj~;X^>sw)Fab-}>j!Q*(mo$CV5 zEB)(&fxm;t@tcvZ3;6l;_W~URg`Z9TIl{o-$K`muo6ixr@4cPp2;9E^9AV(^l)YUchz5TfL9KyuOcsbG(7D_TNY7y_b;V(SMHcU-uk=`(yt( z0{6V$W3hi-(7!I2S+5H?7rm2v0ej0^y{E{$zNd)S_Z0E^o+AEc?>FrVGT=THH{d`=UeO9cq?xlMc?6(G##HSzgWfH0rm#1~Kj!hAs!Uq}TA^My@( z5fvcJ7d7$4RDdvF+{Bkq0m6Jq6JJUN2=k>)d>IuW%$GIs65FN%vV(b!mnS=#8+1V!h8)A zUsDAL*L-UkoZGK$;_IjYVZN@3ucrcp`T8clfeH}j8=Ck=DnOWTY~q`!0AW7F#5Yv| z!hACm-&_R<^DRt#OBEo@hnn~>6(G#FGC0R?YlAc2#>BT(0m6K^(b9Z7tw6Xhj4<)- zRe&(x!NfeR@bjB!;s>YzVSb>AAEW|= z`6R*bLq01e37p6K7wvqGd`28>aBg>qi65#0g!y46ez*z{=I=4__o@J4{yq~wLInu( z_nY{UDnOVYW#S)D0mA%f6F)`;2=il2{5TaL%#Szm$tpmYPciWmRDdu)(Zo+u`TyBF z3wSH)x7+_g5d&MqcH4!D9jIV;x1wTqE4CaE?8NSaf}PlXFLt-sB6jD^H||~E+0T0S z!@S7(pO=uB_4$1LJ!{tNefFL`^UQD0oP&VKj}PomPyvYi#K8U}6@bnYCkJ>w?^6Q% zQ&j*W?;Y3=RRO5Kcsj6uMg<`9X9N4^Q~)A>KCpj51t9Vl1N)a$03v@muzy7bAo5oO z``1(eB7Z%we?tWz@;3wfw^RTke><>$M+G4AcLV$PQ~)9$7TCY90ucEJf&GUn0Fi$b z*ng}75cwy8{iiAbk$)E04_5)`H1TdI4T#<{xy-dax32NVL)CTe`ZexlV|%babO~65UoNm;UIid>=G}ns zD+cz=y8)3`bMN~6o9)5epqhJ6Saa_^^_c+T^Vut~-&+MB@_hpPeN_M=-!HJ=Uj-oY z0|NU4RRAJCD6l_R1t9W60{cT%03tsuus>V{Ao6N1UCpK6D)xr@cZs)E1L`l{3GClh z0f_v)zJU!vhV~b?Y``=eo6? zrcloq8929dP5Yr@tmfI(JUey-4XbO1sc8hXn*0J90kQ9JE&`N&b&b2eJ*s(jb&Y%5 z`r%O=C)=q8#BsBIV84S3K;%0H_SH4+`rO>xqnejjzfbZaJA=f$k>4c&gnv1(e?c_UewPFg`RjrG8!7;izZuvwC$Hw^Y>oOg>n}6{;u<#Bx2tR1Yz^OQs^;Eo3+9fK z>)f07?`zuB%$+yVT)U%U{-?Xk8BKYDh(^ZO<1sb4>>4-OD>)n@|xXH@_qe=e|pUIid>evcavo_Y6Q&Am6!&(Cx3 z_2pmD`ha+Ro_q6p^W1w~`SCgTUPnJ~bl1e!R{nQh6Ys8{KU~`k&}riH0N+jd2#tW& zlK(;@pz%HTUQ<8+2mM??{QMsS`=3+*;&-Kf4)Co1OMqWP&-bfFKwag3(+G&))%rcK z|3d{J@;?K7{r9{^{r^2@u2G)bbk;IJRJYfylY^;mI^@R-2(ga?;e(aw`{0Z z1hk9Zdp>q{me0q|PV)JF&yMnW-`+t!-|yL8KJVMx$>;k$+sboq2hcY1`Tow<^7;PF zR`SEO9-uAdKi3Fo3;DdSZZ4nq)y?F;*Lr|9mCyUaCh~b-*jS!>7<_*~wc@|^dGrCw z7L*5HlU__Up#7Ei)ad*Y)k6af|JX@A-^i@oIYw%=1e zZ@ZrIpX=uU+CzSXMnJpE=WVl_eBL%aua~{ITU#vz z%ij4E3c#{=KBWS%?47q$0G7S;sT6=^@4USNu=9y z&R0|bmc8?p6o6&#d}Re-**jlF0a*6VS5*L(z4O%+fMxHT`5##J&bulA%ij4K3IMVH z*^{&Xt*K>T+57R#FTt{R-c12m_RhO20IRzSaK5%`K%5)b3GCNZ0f>CPz^D>ahH&p?Me6zrQa}|Kdw+QUFQ~`*5tH6G16@bXM3GBC30f>CN zz~~ZFh0FfUS*dML}5cv^-{gEmFkslS< zAFToq`7we0u_^$O_X_NfQvrzl_`v=I6@bW34D3%*0f_!)&*w*fJ6X%XviJUTiUP3g zou8@zEPLm@6@X>$e5e9I?S(!0k$QbkYZ)MNd*{z+86a|d=g(>xAaZ->&uJMTa(m~` zYZ)MNd*?4`86a|d=Pzm*AaZ->FKHPda(m}5YZ)MNd*`oc86a|d=dWrRAaZ->uW1<| za(m~mYZ)MNd*^Rx86a|d=Wl8mAaZ->Z)q7Ia(m}*YZ)MNd*| z?`atza(m~)v2n7aD(LlJze~_agFXac|b>( z9ydhG0Nqslx2gA3#W?c++wEp8yFDO2PS10PMnJvf@6-s0UifH$KUVou8UYsYt{=%W!37asd+1azwWX&M2YBHv#lpp)g#&FI#K>?jet&&&#&cd z`HQt2&{^`AYXo$r{FNF3ogsgXMnI>_H){mM_8A=D`zvSL0%E(|q7hJ&{H+=R^^?z^ z=V|iyYB`|3@{ef*)JHyVcb<>^rUQf2FQ7hG+4&j)u|Hf8;LlNhsYXDT$`8~C=o0xs z8UbA_pI`d``MeKZB%k-83+3<9;{aVCpZB5j<)6@UKvzhk7^$#%lj}s@5A{10W~^3D2)D-_i2{peHwpO zK#lS~jz2e`M&|{E(SP!Oz_Pp_;I9g(QF+_9w5{_#%Htmos8Rl$@p&J{^SLilFaFVl zi<+TN;6*KM>wa1htZ}{V&(RM8YfL|j3-v@`<^IL80fF;NRDhL_X_j3UIM4l#Wmg2w zbN^#mOZ!IdkE~xlPOsB4uv*%e%ExW_IL-Z(=gj>RpL-IXeQU0IG1C8UmPIXX+kCEI z{c*jm@6pcyE1wHkme23_`vPioe^40x=YhcfK^0&5vJM-TCd6ho{ zGJo!Q&cLR+JBZDup)sdz>Acq3{keSJ;%E7@AJbzizwU8;9Ovy>?h)nVigl-mjoaFJ z1CRA*TJ8loYu58aLruNs4%6D>ptWIm8Qc44c^@2A4_h=o>9wu$_rdZ!g!eUj{NVC6--*4qPiXy%ZQUWZ zbxY^TmgX?yyElyL8paIqUjFXaQa)D4_I!@d&G&Sg^`61{yUz)^FO~19jP2_Nc{?;| z1AsNY=SJo4gS@tO@nXu?aa!tu|4cugGM;;EJXia8uCe_c^LjlmSj{S+p{Cx`hiTRb z*0|n}$90e3^Y?X@>-46;DsPv(ZcFDl)}1b%s7WJO<9nY!Q$HB2mhP$Y+C67!?OeRI z^NzetTe@e%y1ri9+Bs|fIag~#fOU-uXs9utE6;U#o;I;!lSXK$ske{Ou|0=q>E2|1 zUH$cfprNMTvxRACzdc>+fi#q*4_gf z)AK*?y?bo5CXHZC=skaY9;fqYqvY#n`0~80eB6%jYt!X(!o+s1dR)(&%IC%M_q@sN z*ylP7G&Jh}A09dLP0u;;^kMg6{FDJTY8Mnn|6y+C`5DWYb9ml?pE{sMdDf$tmOX9=j$tU+P)pVMFqMssgMvQ~+Z8 ztr_6aT0vp-pKgJDcNJi*tpX6o<6l`8I+5wkiOTZx`5auL2PH4uSoSDgcr1 z6xi>q0ucEwf&H#3!0Mp_tld-^Q8u+lyxc`)5@EB7ZKhe_jP3@)rX87gYcve<{GvEnW`rbBR|1JnO$2 z;Ca2T1^79{>j9qi-w5pAQ~`+mt-$_m6@bX!3GClh0jQIBFTk_@umC^1ct60Ce-PMz zr~(lAM}hsvDgcpx64-yL0ucFU0luRc9%v}$ai41hG@BR^;Aa(I1o$b%mx29PDgcpx z9oT=P0ucGPf&F(X0Fi$m;Aas(1bFfv1N)y;0BR?G4)Co1OMstQ{2JhA62Arbw&M4| z{tp#^$o~xRZTOGGa#Sn+TmR4e;k|6$V)3ZIqw_iS30Mu^H?TiV1t9W%fqj1!fXGh| z?9WgEi2TgJ{wx)M$j=V!&rt!0{M^9)JQaY*&kyV`PyvYi!odC_6@bVG1ojuJ07QOC zV1KC!K;)MN_Lr*wM1DnJf29gQUVzWN0H1pSKKBBA?gjYV3-GxY;Bzm)=U#x%y#Swk z0Y3KveC`GK+zarz7vOU*z&8ifDE9(UVzWN0H1pSKKBBA?gjYV z3-GxY;Bzm)(+AuO@VOV@|Ed>|jL%*`@2Pr0y&tgcQPm3?^8&U<)eGwVfNg;$)n34R z8=Jj=-crpA8s8DvzSX>-@$XcyJ*s&@<2wS|V?vo1@JhyJFW~tm=3YQA%ku*K=)NP6 zSKkpD|IU#%$2k3t!1ivOMl_o+6)Q)eEX#z`X?Pzj7a;T|7A76XAYRzQ@rlpYL&SKZ>90pM80ok?&z} zKZ`!op%_)+qP{}w=6e*}tI78(xtDWpz(bh3<@;k?Kjr5G;$BO>U&ZUp_p5jxxrcKv zC*SX(U+4EqmftJeh4zyHNhw^sv00>(9~n`?Y-TXWiD^H)Q>mt`G7) zI@W*aNS^_V@@IfAzpwsU23Glf=6l5RXr@KcltRd-As9@#Sro_Zjl$fEu;- zy8Kgm3^df#d(kjsdVl&-Jr7uysQ_zq-!Iwrc}|HR-L@gm#}_}NrTwT`>wrbS0Bc;| ze{Rk9Io@AB26>I;V~}GW5bOV)b2wf@{yv8v)xM@SS}6DYD~pAVzs+mX8h~~${@cj% z@mkLEvFNN_@%WK{_kk?$YxsQJ;q$$CeBRc0deRidIP&#;vi#ogPX?^+$F{{t<#qEm z^Sa(<`S(-G+2eXHKtB}!ZEUl=ANg_Ji*@STxhd*{HTDTugwOpIpZhDmU!b8_pI!@y zJoj6C{`~N{-{SM{)&`>-+_eApR0_WGO07Tv#Xej(Z zjey8=PsHCCXejbQ8Uc|H4(xAI0f;=$aq#&X8$SO|4nE(9$KM`kD4y>QjeyAS4D9po zum=i?rq z`^i;m98YQ)AoAQ_@bt0n#W->==%r=37aS#@djUT80(|ZT_}mNdxfkGbFTm$sfX}@E zpL+rRoIpeI_X@cekmp{2&%FSjdjUT80(|ZT_yK{2;^*aFK%RR6KKBBA?gjYV3-GxY z;Bzm)=U#x%y#Swk0Y3KveC`GK+zar{frjG!4e!1k|tLF4!H*dA3c zX#9R2+hbDg1-!Sh*$e0`)x4nb`+02NYF^Oz{XDiuH7{sfBV>C_DD#4`-3xfWiMbch z%ksPcKf3P-pIE;m z@Y<{U2=#jje2RRw)pvyP_Z@-l-J0(Rtp8x3VS@dR(1jN#<$G%T$me@%_yK{2;@?!u z_tePqJvDs3r-r{Q&`>-+-%}&cb7Xw}zKqY`m+|?Y8a^Kj_WeHowcsp0c^5ufj=;q!HLe7>iK z&-c{ud5(1T>rc=NbXcDnCLapljs6&I@M(Ci->i+9^18TprP2#Q)mQ4K4oCvP6Z(HsRH}V85seK;(-B_KT|kbhIeL*X~o2&g7D3p5mdbB%y%Vv9gS;kVQXs3x`wG!%Yojeu%mn?OV1x77%!CbkPS z6n=Y+fNElgKttho)Cj01b_z5UerJt_K{Gl2F)x=?ehQc4N5l~GW5ojphhXy1vcG z{}ypg%h<+%YNA)5q439P1XL5p2O0{0f<{0!ablpM@F!^m#OpXY!1Fp!3G7c*0f@YJ zU_Vp^Ahzq%f&DWo0Fgf%*gvNN5c%_g{R=7pk-r$&zoY^X`OAU*D=GkyzZ%%TrUDT8 z>w*0nDgcqc8Q8z20ucGzf&Duw0Fl2N*uSR&5c#mc{(TjI$Ug||KU4vT{G-7BV-vDgcpx8`yuR0ucH4 zf&C9E0FnO~*#D#g5Pjz70MGir1opqG0Q4{7w*bGV_&u=yLj@r6KLdPE{-e2e+hVNy zo~qqTIUWHux8c7{y}K|-7Y5Y4&Wi&30V)8IUmVz9q5=^4rGforDgX`C&%RtEAlAPk zu)k6TAo8mM`>Ry|BEKfEzg7hx^6LWo>s0_EZw@pRexOD`v9W0ucFaf&J|&0FmDj*x#uF5cyq!{oN`6k>3;8->U);`H;Z=J{5q- z?+@%BPyvYi!NC3@6@bVe4(uOM0f_w3!2U56fXE*Y?4M8pi2TXG{wWoJ$T=3e7o#d% z)CX(q6R-&1H?TiV1t9W%fqj1!fXGh|?9WgEi2TgJ{wx)M$j=V!&rt!0{M^9)JQaY* z&kyV`PyvWM_X2$G1^CUVzWN0H1pSKKBBA?gjYV3-GxY;Bzm)(+2?I zb1%T>UVv{7G!*OSUO=9E0Y3Kv{GdQX@p$?GAoAP`@VOV@b1%T>UVzWN0DpU+p?Ln> z3&?XXz~^3o&%FSjdjUT80(|ZT_}mNdxfkGbFTm$sfX}@EpL+p5_X2$G1^C*;^XKaYjn6Hy#SxD7vS^t0(`z+fX~+p@VOV@xmExOpRX6-^LGS%bD*JE zf1pM{jn6Hy#SxD7vOIXG!)ODuNRQ#UVzWn z3-I}R0X~06z~}1)_#uIY;^*b-1?2g9K?faMd0v3e*9%xbUoXJt>jn6Hy#SxD7vS^t z0(`z+fY0*+eEyDr&%FR&^@3s}6rb&c+zZ&3s$Nj<2Yh1eyAx6`;Pc77fY+CM0lw-5 z+6WTcqM8>p?kTc8#$;Z=zSo+00qa-ug8JOR`{THr7gW8V{yl+h&A!L?13*TQF)UO%17gpB`>emdsJtn8?1so66cLYB9*3S!=SB&o3A$h)bh_9|4 z)~_FWdso*E8^3qRvG{jhJFI#^^&NqG35b26nioty^8!Bav3c!~bKU6fBar9&2>5&- z0iW+9;74^Ifw>D(TtLlxX&cy2p#l*3l!1La6@bX63hdjf07O1@U_XrtK;+W~_S2~V zMBX8=pI!wZ@)-jA8C3uxpDD1PSp^{SSpxf6RRAKNEwJyX0ucG^fqf?xfXL?v?B`Sg zh{=T!lSe7?YbeieYoYZTHy^LMiUuJslOcnptUFtA@p1t9W;1N%i( z03u&BuwP6CAo9fn{Lx~G08hSTV84_KK;)eR`=wO?B3~x3UseSm@-Bh>aw-6kFCW;i zpaKy2ih=z~DgcqM9N4d-0ucGCf&FSK0FkdA*mqR{h}{P*mqL_h`f8C zq3~;K1XL621R4szu0}vLv0k8|@at;?R1+Hn8VbLmMnE;OQJ|sl8*2nq6PpAY3csmF zKsB*hprP=aYXnpiTLc;kzokY%HL+Eoq3~O41XL5-1R4sztwumKv0b2{@Y`zyR1-S{ z8VbLoMnE;OQ=p;nJ8J|~6T1W&3cssHKzx2Z0{h)m03zQ#u-`)kAo8AqV;mhTdItX6 zE_*Gm{!Uig;K+dOQ4?%KKsCWJ2&g8G2{aU*Z49U;dIcH^f1E}@HF12Pq3|bY1XL3z z1{w-~l14zhj*|mCuk)0^{!|r!$a@F&LsbAeSUerrKcfN=`Llujb1DFlKOfk?paKy2 zi-G-1Dgcqc9N5320ucGDf&FVL0Fl2Q*uS9y5c!*d{aY#kk-r_-zoP;W`MZJrdny2t z4-4$yR{@CpgTVen6@bV;3hX~t0f_vQ!2VMefXF`!?1!rWbfEY=!0#hQ1omI307U*} zVE>g0K;&Ns_TQ)gME-4H|D6gz;Dqi|EdDezlh%g z{GQ_X!2S;vfXM$0@ICpD{@!n%V*Jl@sU@@=5XV~Qz|sWX z^pN}`8UZ~h|F}j#56C~E5zzheJa-2M>EeKz=e|NCpr_@p)Cg#(e10GQU4EdJ1Nv|B zH);g*l>8u#fS!~etP#)?@;7S)^tk-38UZ~ff4fFNkILVz5zr&@yq75$q^5xSIc0q{ z0(w^dG>w3siN|_7G-(;2mz4F<22lTOg-UmLCzeURdeJFpMMnE6P->DJM`|^1o z7$$#@mIHcE{$7oM*r)TheOI1sH9dpW98mwPY@kL!Kgs9U`lEdQ+i=W-35fEAV^D2KH`IzaY_47H%&+YxJrOzhsvwrMUybvk(7Jg3oJ?7rR zV{@Nn{eNZ~_;W7rV|gDK)ob*&8>pWHh}ZC7U|&AY9#a5__46_6_wUzqnQ}nC6bl>w z&PqP^e^Zvv%b)RFZR-yo zX&>pVET6;hY;Uid&oN|qe=Dz>_YtzZ4|v_;>Gc!(1^l#fmiJ+2-VV9nkmde@&;12I zs^{|llKU#_=KTSmKUaMIT=99EWPT6%yXRuchG{)O|3`j=MnDV6=kG%Q zJD#h1k=FN@rl{%LeF7HY`v&&8pOELz0dKs1?l)xl*u>}01D|^gK7T&={Q2N>FU9Bn zk3T=qP`tJaGy)>eb0T~`PVfT)4aNEwYXn4|&nfsz0}Vxm8^O5cv&({f#OBkq-*&2de-?ep6t7vkE}u%m)GC zZw>78d6_()m+|?$jL+v~{9S>D;`#HOl045T@p&GH9};LN9-sRcdF~JR+zavf9D>j1 z5PY7i;d4*HKN@H#etzyTk}z~^3o z&%FT8^Vfey{W#6NfIRmCeC`GK+zarz7vOU*z#ElCoqGXU?gjYV3-GxY;Bzm)=U#x% zy#Swk0Y3KveC`GK+zarz7vOU*z~^3o&%FSjdjUT80(|ZT_}mNd&4Gqu`{iCho_hg4 z_X2$G1^B^%hT{2hFCfpo0H1pSKKBBA?gjYV3-EUY8j7EndjWaw1^C zUVzWN0H1pSKKBBA?gjYV3-GxY;Bzm)=U#x%y#Swk0Y3KveC`GK+zar3)eFeRXD{G5 zta?G?`W)N8>IIGKg=~+i7u5TKHbv3vCe>cRXFWE10llS~7u4ql-gl~bL49uE?NQAO z>T?5cj|pX7Ft&RE&o?pm0(x1V7vM+t9f7?1j?nmbi}?J<<#z*)A=PYr)ZprLsE`TH_?zNd!I*Us?w1R9FR=Q%QYzNd!I=S6(Jr-sk> z)bRPdh|l-b@cEt^{*gdK@p|$-HS&BN9iNW{e7>iKe=^WeJYPO9l5<|1rx-{6I|2Ia z5&<#3b6~%W3P9wm2lic603u&2us5Dt42T@%kx>bWsnvO`1?6}QqSb+S***x zhCYz}OE{h=#i1hj$d?W5yQly}zFc6x zyb3_%D+Kl{ssKd3QeeNb3P9wm1oo?{07Sl8VBb>TT|+AZV*9Nb;Ca55`f_)z1Bmt4 z4)E*?>jd`essKd3USPkz3P9u=1oj)M07Sk~V85{nK;)YQ_M55zM7~*IzqtxPQ{Xw?FA2{wbiUX#z_Z?7f&Jbp0Fmz#*zcAFl!s`3Zsji7Eh*pA_KF6DJ3FKJQZk`%_f_BJUm84^;uE zzj!*Ze?|o$@@E74=Trb9e?G8(K?NZ47X$m3Q~)A>Ik1041t9WQ1N+xh03v@quzy1Z zAo4c@`?pj8B7Zxue@6u%@^=IK_f!BP9~RiZuL2PH2Z8;EDgcpx6xe^P0ucEpf&Hf{ z0Fi$d*bi3$=rr+pfM>rL5!ipB0ucF^f&Eu10Fi$k*ngt}5c#)({dXz=k$)f9|DXa8 z`HzA9PbvVhFaI3iS^t;7{#O-%R@U+Qn?^t!OTP#9f2aUN{%3%%ABT;jL}hRRAJy4m1>gphiICHw5-KssKbjD6k)_0ucF4f&I-Y z0FmDk*x#xG5czF^{p~6Mk>3&6->Cu+`CWni-6{Z)-xJv1s{#=Dkih;v6@bX^59}XM z0f_v;!2Tf>fXE*X>>p79i2TvO{xKDR$R7{vpHKma{K>%nDHVXoITpGXqbgj~{WbOp zScLBz*q^2X5P84Ao@Ia-PY>+RPyvYi%)tIE6@bXk4(!iS0f_wE!2UcHfXL4e>@QFO zi2TC9KKBCh+zarz7vOU*z~^3o&%FSjdjUT80(|ZT_}mNdxfkGbFTm3W0O79->~k+5 zZw{R2UO=9E0Y3KveC`GK!GVV2=jUEPo_hg4_X2$G1^CUVzWN0H1pS zKKBBA?gjYV3-GxY;Bzm)=U#x%y#Swk0Y3KveC`GK+zarz7vOU*m_p|o`T(G|^0^nV ze(nYMsuvU+VY2H5Q)=I+dO>4ez_BnX^#a~c?gjY&YAZCT9=>T?5a zilTQ`^Mb~A1on>!Xjh-_dI3IPFTm&P1^9fu0H1pSK3^}u=j#Rde7yjluNUC+^#XjpUVzWn3-I}R z0X|jn71f!EL13&`{J0(`z+ zfX~+p@cDWHK3^}u-w|jiUVrWd6IrK%Uy z`vIRA`|gC)3;29;FW~j%Ylrx%7ic3$Y>R4M(730__860S0sCHS<^`-@%?s*t1MiRH za$ZpNg8KIawl(`+^&O%9J;D24^&O${Z*g<1Ov>L8s$S6O1#F9|7gW7~YX*q@VRF7+ z!27ShBd{m5{yPHmiqTy=B+u6l@oaBEjMcTn>e?aK4-xyr-+AqjpI`L?K8gBgQC%+> zd)Eutw_0<(fc2~E1rzV}0?u`#yN^Jg?<3&zeFS{IkANT5eFW}pAjQ?+(H2{x-n(W| z4TxiD*1&!?6@bV)2KKY707Tv?u%AN(Ao4i_`?*vABA+|3pGO5C@_7UM`BVTRpFgm# zQAq#H-_fq7onry53y9}iFerQfEyzL&0I~kUf&C&X0Ff^m*e|965c%Q(o;lwV0iJxx zzBqz^D{chBuz~~fHhg`LTh0FBO2uj|=RNR{@Cpguwnp z6@bW33h?xUlLI`T_bGw>sVV@G_YUlbssPkqJRR6SqXH25vw{6{Dgco`AK1U30ucF& zf&EJ=0Fl2O*uSCz5c#Wt{c9=!k-r|;zo7yU`I~|LTPgsNza7}WqXH25yMg_CDgcoW z3+&%l0f_vA!2UxOfXF`z>_1iki2Rel{!+fji zsXyI+np-WQ<$yTWItTX4r~pL1dSKsG1t9XZ0{d<%0Fn0y?90b~KDJx?Zx?lFJNkdS zXgcM2-yYxpcG0x@`FUS$>EAAzM(gK&!T)wqdp4eQVL&~r*1bq0pjqSxXaqE~{AC&e z%_M)dMnE&lU#k(&4D#1$1T?)ozYErmLApPno_ zK=a7osu9rK^0#XQG?)C{8Uf8I&wH7IL23%97gg3*BcMg(Ptyo!Vfnlrns5RwsjQDi zKugH?)d*;Dd5$kYi^=EDb9wpewH(lL@;7M&)I~nO?`7rh({ez|$UmSF(9-hxGw&>) zKl7#Jc`vgtNcpv|sI2_HnHK zK9<*)e^ToKT2KBdjeypb=e_xO$j{Bbn)e}m{#o^XdDQ=pW3s$Y;rj*}isK;fW90n; z=lxXxB0oK_KSKo|@-qYbvs3^gKRd8LM+G2qW6%B$h>Y#jrWpV8?~Rj98Cl*Jd0hT{ z@plBQj%Pmwgy*^17vt#fiLtI>b?WkK&c_{jK4;+bafjz;w<|_fxTwj>pI7&H$NuO;^ueopy4=H9|%bDw4Xe`Xu_b1v^=Y#(Q%dX3(81NCzN z@fscs?90d5V+sJVem+M1{{3_FJYmZ?ev3hm%g6rKTA0t{TgAG5p1d8lDP>K1Eov)<;EPLnO6o6&#yoUm??46&Y04#gwLlpp8N!XJg ztZn_OmH}E(*pnZm{535Dw1TiFKTvu8te01APkw;%_p}Vqa>AZ`f8}}K=%U=7d_U#y zYZ;(rg+2Mc%0JLDKzw%g{BwCg?3?!Fdng~FWq=kJ_T;-O|3b?EEhg;AcT@homH}E+*pv5Ao{z0X zl-rZe8k@GDKM{uxH&%wS?8L~pPsKdMc=ztiI_YVhNiR1KEBbIvbd**oX_1eU#X&TnAZ zJ6}ctSoY4BRRET~^DYX&vUkq;9V~n2%n!h_ch39*EPLn7Pr$NwzLEm4?47Ty04#gw zt0(}=-ubEuz_NF~ngX!woqO)++}?Rt1z_1bUqb;P_CI@a_P;f?3@m#;-g8gq_RhO0 z0L$KacLiW|R{_q~Rt<=A<2r%;x+(yXuNT;_uL2PH27&#CDgcph6xeU90ucEof&Hc` z0FiGN*l(@^5cw8?{gx^Kk#7~)Z><6l`8I+5wkiOTZx`5auL2PH4uSoSDgcr16xi>q z0ucEwf&H#3z_Rz}*FynV_Re=x0G7S;-4%dk?|cshVA(t0Qvo1e|6YOp-YNi*?-SVX zs{#=Deu4e|Dgcom5ZE860ucE@f&IZM0FfUO*dM9_5cy$&{oyJAkslG*AE^Qm`B8!W z(JBCu9~0Ogs{#;tufYB|6@bW(5A07+0f_v>!2To^fapK=e17zwleG*idw;&CC;-de z`Kb!PvUlEF0a*6VhbjQnUf7c#sn_?kmH{HScm9l)0V20|{;ZY(BDZ(`oR$G1w|D-$ zmH{HScm9Hw0V20|{-TxvBDZ(`l9mA?w|D-smH{HScm9f&0V20|{;HM%BDZ(`nw9|~ zw|D-!mH{HScm9T!0V20|{-%}zBDZ(`mX-k`w|D-wmH{HScm9r+0V20|{;rk*BDZ(` zo|XY3w|72F%K(wvJAYrx0Fm1}|3J$Ck=r}}P|E<3+dKbA%K(wvJO5bA0Fm1}|3u3G zk=r}}RLcO7+dKbE%K(wvJ0GrPfTj||1M(wu{C+MEXi8yEewgwRS_X*R-uV|=28i6= z`IlM-h}_=!S6T*$+}`=uS_X*R-uX9L28i6=`L|jIh}_=!cUlIB+}`>3S_X*R-uVw& z28i6=`Hxx#h}_=!Pg(}3ov2|P?Po12nau^Z;K1{J^}H5FVYAIKR_cO{KXmp;V;n$2!E+YK={iv z0>WRe5fJ_gjezi1Y6OJ8N+Tfr)fxfeuh9qyf2~G9`0F$R!e6fu5WZO>ApAg$fbchH z1cbj)BOv@BjezijH3Gulq!AGQW{rUGw`c@}zf~h3{B0Tm;cwRn2!Dr0K=?Z~0>a;= z5fJ`vjezj?Xat16S0f<&5RHKF_h|%#zh5ID`~w;R;UClp2>*~qK=_9>0>VF{5fJ`S zjezivX#|9STq7X-6B+^GpVSBl|CB~RczQt>28rVU5LtCBRL8>Pb}V!+_WqI2i+um= zqve3``MikFa|Ar|13>ux8Uf++908x_2>3ilz~?yvKF<;Gd5(b3a|C>zBjEEK0iWjx z_}mNdxfkH`c@dw_i}-wA#OL!OKA#uy`MikF=S6%zFXHog5ueYC_ypBM4@ zyohhs=L85pP$M9GJ}=_)c@dw_i}=BM93Xr?FXHog5ueYC_ypBM4@yok@| zMSMOl;`4bCe~(@#AUu7*y#Sxji}-wA#OL!OKA#uy`MikF=S6%zFXHog5ueYC_ypBM4@yok@|MLg%l8iUjnP?Po12nb*Gg2o(~?J=qKg1j#@>oW($dmX3|5WeaQ zjrlU$V^Zu3-HT6q3ilz~?yvKF<;Gd5(b3eF2~62zaBisPi0wEYA_}d5(b3a|C>z zBjEEK0iWjx_&i6z=Q#pC&k^v=frjGs(gy&M=U#x%a|C>zBjA}I0K(@v0zS_X@Oh4a z&vOKPo+IG%908x_2>3ilz~?yvKF<;Gd5(b3a|C>zBj9r{z~?yv{((S4@&5B1fjrL< z@Oh4a&vOKPo+IG%908x_2>3ilz~?yvKF<;GXeISH#1Q33p6idS+cqe|_RiZW0L$L_R0_bdcivtBSoY4RRsfd0^Jx@-W$%1i1z_1bpH2Z-_Rc#f0L$L_ z^a{YTcRqsxu)73&O0dp%ij4M3c#{=KBoe(?48f004#gwb1MML-uXNVz_NEfuL7{_ozJHLEPLnk zD*(&hc`X9Gf4|P_a{qh*8L;gA_yrY!WncY1U44rzsR6vbmQoFtz4Oirz_NF~v;wf~ zoiC#REPLn6Dgevgc^3s>**jlO0a*6VmsbFmz4H|mfMxG|MFn8lJ6}lwSoY3WRsfd0 z^HmgpW$%1d1z_1bUrhm6_Rd#V0G7S;t_r}icfN)KK%9&0$*bR|tN$+D#>LjDpO-gL z4Ty8|rh)xtDgcph9@uZ80ucF@f&Er00FiGU*l(i(5c#%&{dOt0FfUS*dML}5cv^-{gEmFkslS zGC<_^&R^6rK;-t$U(zx_A)-pik_Re3?GC<_^&R^9sK;-t$U(+%`A*D^rl z_RinXGC<_^&fnBBK;-t$-_kNbo zX&E4Ld*|A8HvOa(m|=X&E4Ld*>f(86a|d=bva9AaZ->pK2K( za(m~WX&E4Ld*{Qo3=r1>h6m);@9Q;wKTn&bxbE_uMqvB!yAj{32E?@(d*?rB86a|d z=Ray0AaZ->KWP~tuI1R1PcFZ&$NNTS1clLm&J66&QUTW4D!@8N1vJ$7HyduyLTIS* z_xzV@Ay`+a0P9K>U|ppG8fyI884qb8SP!cJ>k$=TJ*on%$5cQ=jrZ_?7D7XfpL>rM zf_1M7u!g7r>pm4=-LC=~YW&<=wGbL=>P;^}WBUFc*K>gNgbHY=srNKt`fCL1bQNH= zwEvu^6~Ve%1z6Xr0P8vxU|p{Q8ft8ZCM^W3j|#B*ssQUW6=0RO|ENCewy|=PMzG31 zOX+|7TNUN~sl2_P)RTbqlnQ95@o&w(qM{@1J(!5XLntQ%B7Lyhgw zq=jIO?d?)NHrWnS#1Cs}n~v_W<^8wZFUtEfKX<11xvhPV&3YVIEp6W$wIW!9RDd;D z1z0z!0PAKIVBMkuEP60lx2XW@b`@aVp#mCe>fI?!Yqv#fKeO@OE6&&Z0_y@5U|pyJ ztcz5DH9!Se7pnm45*1)wssgObRDgB43TUX&%SN?-yGNAgFfH}Smfm+uJ*Rxm<9KNk zFS@+V@-f_`b--$Af8w$2V&&FuJLY0wmFHvmHQ%X~prOXOu>81|woSPo^K*P2Y3)3& zwS8+;=R5!Gd-UeO8q>L6d9LxGo){Wx>OF0k@;))XkE8OwFskk3dBOF1&0vjcUGM9o zJ6`YB6F@_w{<|9TmW~7Fd|*w=eRVUv$=p}*W4o`G_lGe(#>@S1Y|p_a^f^Mgf8M3f z46O38Fur@yLwX@#J*)z(zoRdAj33_GW52&12UdArkmpbwJ7)%7zoqY6<9hB;?yt>y zVPLiN-LpJ*A60)|G@fis&pTr>=V06CzJed!+<|;z>m#i>j(LAvI|6HT*DqMVJg1lt zk7Lgp@|dNXIK^KU zXefSOo~M!L^FKZxr}%uF;`22Ze4gXu^KpyM=YM?edHCi)L-G0sY6L{i^%p?+eE!D| z3N#e?V2yytZwl;hRso2d>jQxB`5FN}Un9WhYXtZ^0u9CE^ZB1Vpa1ds{EyFl1E0T( z;D-bnis#RLgFN>QeEu$i&wT@*`vyLr|MB_!kAF1KQ2hLSP9o3W%kcS}gnu&7P&_`L zqsZw6?TT^a_wsO?kzdxA~8hKs%SPC>0ehQ6%$fpeK*>Ak>knDEgK3(9PeaZV1`SgMF8B_owpE0nXNd+MCnFIS-Q~)BM zHL#ye1t9W{f&J_%0FieJ?B`GchX%&FTmkI2bRRM^+OJKj83P9w`2lgwd07SlGV84M80ZZznTg_ zkYcO>7ouDE#Ic0oBA7fri3wsS!|3Y!zrI z{MH%))xF5l~HR7icK__8I}z#14Um!tbaNP)+O$afFy_fP?dyrHL+iyq44`_1XL3T1R4r|phiG7aZsS4@CR!IR1=2;8VY}?MnE-jSfHWshie2> z6GsFZ3V)kkX??3?cgc=8Vd`wvwBBL66`|5ybe@=pT$PgMXS|17|B>Uj_of4`>#|0BL6zD|3(EM@^1tC?^FOH|31LeD}D&@AslTa5MNG0J}&riC>bKpf|++r1c7 z;iBeq>JzXC-#4&7O$8wGet~^|6@bW35A4rS0f_v}!2T>1fXL4d?9WjFi2U5Z{yY_c z$j=Y#FHix9{KCNgA{Bth2L$#Ps{llPNnn4e3P9wS1@@P#07QO8V1K0wK;%~i_E)O_ zM1D$ZrVjZ&U$@d{AINSOp;Rn*#frRRAKtC9uC$ z1t9X<0{h!l03yF5u)k9UAo9Bc`@2;DBEKiFzgGny@*#o!eJTKv-yhgNpaKy2gMs}+ zDgco`9N0gi0ucG5f&F7D0Fgf)*gv5H(BV4vpVSD5^`8pt|Ed>sXftvzI7$!5y#Swk z0Y3KvJbeHVKKBBA?gjYV3-GxY;Bzm)pA%>(oJof^8?gjYV3-GxY;Bzm) z4+%6BKQH$J^4tsXxfkGbFTm$sfX}@EpL+p5_X2$G1^CUV#6rUO?6- zU@iQ(>;?3msuwhVZ;S0+^@7H{fbCKBg2vxIustT#Uch@Bo4tVEQq2n*-x1ip)x4nb zds=LdYF^OzI~KObgfcG}+r5D2o0xk6y)4fQ@T2>VKwf=EX#8CVpZ~c0j==VA&36RW zuf8MDZ|dJMs_O-fzl&jCt*#d|t{t*H{?6+K9IKP^^#a~szFvTzP_GxTFO1Fa2t41! zJTKt&=Xn8sbn^o8YF;q@<^^o;*31i7|L>d^lzRcMC-(yKiS;`Iuf4jD(72~q^@5S_ zDY8Ar<#z<0e`4wdyxu%7z~^3orx&y<#*y!dwD$M9kI?5C~`Q8(s!Bq84G19r8bz5`ako8--KFH6~YK8JWzYmLrjclZA_WHc?{ULtpkXD*&Jvne10$FHw4aGd*7YcHmZG?=k6Q$x${0u z-al|Yx@$7Mv{K%udHh*{AD{Pe@-h9Ljh4RGw6v}BKFaH9=^lFChgqL(yGSvr!bMHi zCtwl2rEQ&$0oMQj?VSmjT}740zl07!2n3XXf=X0~8?LiC&H}UQ?Bnb+?z^+k=ovPcxF1Dw9SKBXi+`GU3y-Z!F?oSyn_n&b0NpR-YaZS4K|83I#d>6)8 z%qHe-wt-oeT^8dyY|JM@*A{oR*_`>4@A6oH{?t9GtuEfRI@-0r`L)^fJ;$r=C)g#& z-D3u9T)!hzHe*_kZz${W4ZQ2$g4KP?$G)Boy8d07w|kG;1`ygg?z5}gv%b5oLCPM~ zw50oeEzjnd))>HbeZ3zb^cUUnP1=2b-`Dr&#l9a`f9NrKj$N+5QOCm16TG|b<;;`x z3vJ&LD0@Q8sh(1R@ZPSzc7KZ*{mhF!7TX-Hb8wgsTgPBaHV4SlS3m#5w*-_l`f3|L z6I1@2K$iFW0`&g(=A7O9{s2DH*$REv7^1Id07Cs+1If230}%8ze%`L%0D^upkbH+S z0738f8Sv^K{JsMHo`90(-$vH`XQq7=d%Q8g^Sp0g{QK^sXq==MAWvJn`k9S-1|U4Y z*!<2GwHY9^A0K>5lP=k2R6n%Xd~}WJ`bDrT{HD{_^8n#p3%w_}uc+Tt#(l5vcIc_^0KpFP9XHl94{ck~w5HwvRv$PUe>@TI zZu6(JdDqH#SDWuXU*!OTjYmCaVWY#Fw+-fSKl1Z2kLTJ)(6_v2xlsKA2z{$@otJIf z<0yMr^Csn#0%RfP{F$m3Ak^RXO_jE1%+fU$?;3YdSK7brUYW9aU)wci^?J^e4FG~Y z>UyDU#&b2+bzM)-0)+RV%QZUg7@{qH^LEv*_hlbu;y2yV?K{+APk1^OnIF^dgUK*9Z6sz9Ep@&$;M*oxpn-^ z>pohXcdsIm^aAA1mCpUlgT*)KcYr)doVE?xr~l@#xXp)0?C;-Id{+>97z32WdMTykU(c%I48bEgIn zI58MN;710MAEgXH(BD0f{2s~x1pU#0N3J$&XbAAn5NINPaJ60D}IwK=R|2 z0SNkg2a?}M8GxX_Zy@;z$^ZoYS~3#;`@rmJYatDgzMoCkK*0Oc{Wne|RAIBa{IM z`bP$mKS~*Zpnr5A`FAJ-5cJ;}Nd8^Q00jMa2alwKcfsl(En^8`Ohf>5cEGENd5$60D}G(0?B_- z8GxYwr9kpuRt6yGe`dmg8nxG$)BhUKp4L#1(H8m8GxXFN+9`D zl>rF)rwP=9oFzO>fV19T9-QBhmkLi0z$y2PK=NlQ0}%Aj3M7BFG5|sUoIvvDDgzMo z&kH2~O=SRr{i@E`QIo55cGc= zNPev{0C}PCwg7yS@b*CRcPIl8^zRHLf0r@rF)PX&^HS{Z<#za;?QD10WMq={*7RSJ*|!fgTgdf~GH zIP>Onf#jc81|aCa5J>(-WdMTyOM&EHRt6yGZx6tG!dC)H8vR$50)+Bk3&5GnUk||1 z-w}Yb7VZqdSu5WNz*$G%3?%=SG5|sU?Esv$NUXiQNvFqJuh5Tc^*cab*8K1A|DHJG zgslMrXDk2$KP`})u>=VEL~`aDK+rKa0Rm@j00hq30SKJ61`s%F5g>4U2SDKX9)Q5v zI{*SdE0BD%G5|r3eFiof@Gas1LI3xGz3{g&|eWq{w8Grg8s@t@~e~q2>Le%lD|b6fS|uR0DtB3>31^duL+=s z<2O!i(t)4rMgDQiruzcc7=8dCaQ6jp_XTj?Kd56+-hBaj_XTkG1#tHTaQ6jp_XY5T zvQ6H70Xp{uaQ6jp_XTkG1#tHTaQ6jp_XTkG1#tHTaQ6jp_XTkG1#tHTaQ6jp_XTkG z1#tHTaQ6l9en3fMKlcUb@dE&XyDxydFMv-5lr-ht7oc}v0C!&icV7T^UjTPs0N)i* z($wd^0KNMHxcdUQ`vSQ80=WAExcdUQ`vSQ80=WAExcdUQ`vSQ80=WAExcdUQ`vSQ8 z0{B}4N*epPFF^nQ`2uvr1OS20r7ytuEct?|xBz=E`GSEz$oaSA3#Q@%=Ewf(3+UT8 zeF45@DJ~eq2Dz4&;)1C&0`qMtE|~gv4%lN~i3`T<3wYmR_67Jdj|<>4oe|J4oe`$) z3mE^oJR@N5E@uSFFP#zaH-j_A(tW|y-63;z>AqlaXOQz_A@2)VtNZhP0sZs)0{Fgq zU$7JxEZrS0q%UCp9Z0@_{<<%~H!RjO0&QRVjxhL^z$h}dOJ{`nJ0oE4E@uSFFNQP1 z3bhmQZz5g)&ihr%r+>qk>F>PHQ`z;Z4jH6aqkr?=8oD(Uk*p7pf3$5{?yE>7{xQ`?umZiTfP7LfpTtzEPaM+>Jz1*H`@xf04=nSwx>;qNDWPV@W=1NpqqarBAQcfF6D>AsV9x7pg~DCOI} zq5B-BJhnZlNtJ^|9-aFOc$=+#4N$&|?J3{Y_6r^N?(csuQ`f2cQ^w2vC)|A!oVj&e zlkWR}o3uZle;EtE3xf~y1I;!t%d*R2T!)SMMCjV$t~Q%9fAU=(3()?$C$-hZyH-cL z_BX#ao4)6G)%^s!SWXfht>+ua`J-&f={adiQZ~55Qvq9ItOY?SnoK8E( zeRg$w)_2!6NZEtb2Bd4;z_U4~H3o29U+)G8{Y7_tlXlgRuW+uxu2nV9lj{@$E={r&(x)7c7r*BGL&X8=O`wg!@KQwAXDYy7-jzX1gOWFYwt zWdMTS?=#@L14qXLKUMDu+G34*hTNq#1BCYbeK@?HQ+De& zfS{krFQV__7b(9O{NggT6REbqFJ#PkrSt%yUw$VHujhz*CaZh@HF^dhJnwal`ftH% z-8BxOu49@yr+>Gg9gE$!$qP z7`ObU(^sDX!kk*@J;8lN{iZVRdwsVr$#I#*%S^8*xzk9y9+Mu+np4$000{P| z>xHrz&(&Dhbv-=`5Z;3>*XX!o2)|pQuIdB4U7Mr_NVe_ zlKR}Zs$*;hb7e}m zqA5B3@ArAm{RQ0TBHZsP;qH6k^d;xH$MER9Z*aFmKAS$hmi04hmTjt2jUIpW&jI9p;(q^hqqy(gH;DWEdat<8y&g*L}+o}U-@c<^)Le%Ai1xcB=uai4p)iu)S- zjJS{aE#f}rpBAStcPEi}+rKDX?YDm;{DSlzfBciUkNLO7eXshKxX*!aihE4-4RP=L zo#K792_Scf`yBYXxX0~Z6Zh|%UlsQ|#jl9_9JpQF_l_@%``*EvcH4eQ99s=LqbxgU z`h4a;676mMA3f{i_TS<@U;j(o=PTppZTU}e_oe?K?)%Wci~G3#o4Bv#e-(HC`7h$` zBmY^PzNy)b_`CUD=W__|b-{Ty?^}8>0YT?;3cewrq|y5vL%%VQ`c29J1pVoO#`=D{+%~<6QT3Dh|b3szAHd+xX;ym*GiO< zwm!7A+1DNK_B{jc>n`hFLB*0@fIPZ7UUfWt%~0O=pnP}Ev$lTuoXuzR%;>%aUG1a$ z7WD3C;j^(pj&q%3J_m-iF>SY1^#FuAJ(h&mb@pcc1`zbVM&Z=Erb(x-tLN)CUsrG+ zC%CU8xQ`p$`v~`S1oyFm`#OUA9EAJ2fxk4Mr1|dxpQRKa=zTolKAv!&qwsz}N%Oqh z4!!$7xW_VZw=Epou4~fi`x-ir+2C6O6o-2(1ZS)U<2BHGOo#s5fRg5UkDbu_J_>KM zjn8e$``m`N+0N%F<+lZtH0`@sDL~L~4Y;J+#Uuw~Q# zKLGguCc0L7fIL|II;8-4sQC3t0rC*>8^DRr}cs48NkL(pKBlIu>$pt>hm`19RT5d?q^oXR$KHNK;YHSyj;Hngy-=&0D;f7-|;<``WBlX z^78}je@#G1W4}w30tCIsaBz>|;2y)lFAXSZp1({fK+x|FBtH=S;2JtF>ALp2taaAT zQam%TGxN^<417y~67cG0M)51y-s2hgOz{l*#pXx4tYiAuC7z-DVu)u}FnJ`So+IiW zOS~WxG_pWOxE*?GBS5h4`atrwc(mI2BIIn*{^+w4;Bh^YIZL3PvaX@}W#rf1vkX6k&V-Xo*lu4CWVa{!^8TLa0r zDFYDnfHV%`yD0RV-2|b4Y>PQxZhL2{kseNMFA!M zcmK%m2&u=%aQ7|nvjR$*^1go1dpr$y`@{PICC&35o1*u<9e#B{NmJg(4?VVB*QC?u zB|0A$xQ`3m?JX8E8r2u(`_{~ZI@^tZAlmg^w; zBz~t-fILy0b#^L=q!%EMoz@4EZ%_sx=&MiLsNVpBey)9-$5Ygo$ojkbyv=$LAiT@{ z%;RJ$_cQS7XMF$T`LhE`n)l7M-|?7+`WKrY@-q?b_q!?hB>^RkJzuL7An11mlD|$F zfS|uLko+=b0D^vZAo+pl2iMSfN!PXCWv#Pzmg1R#otby;XW;H<;MLEJ;#aV}$20Jm z;u-Xd&5v|h$MmmDJVW`#5YMa_Os|p85p|CxUXTeI*{1&m2=f%(M3aslQ}eFu!?@!` z@55fT7a-Y(thdWN;@P%&Yv*}FFw%yIt`Aoe>HoffL$G)Be2sYXpNX~u&5cK1|>lpVL16v&1w6?n~x{ld= z`--+Z!p-V8Kp3m8aqOt)Wo)0mY}ouxAJNrEw)MNs=hQvV+gWXfF4u5RHU!9Q{0Glw ztTM`(fV%ul0QWZ6b2Pe%rp#TRYqL%D$Gj`|k#Wa4)o%^k@3QXQ>JLEZ|6HH_>Gz>c zZKL0BkGhWPefKasWn6HP+6)lvG^}i{$#bO#NS#yl?ge{f9qvb-tC#tC7*5%&V^Rzt z{98ZQ`F+vyW|2)j=XcutyHTDym_*{SIse9#J0^e}toQ6v3Xm1zzCQRjcvtCnfJ{j5 zHae)Z!~422bZZ8WPKzIokZyy12gu>#rzr)MvFXAn3iXa35DVZClf%(_u1p!Yjo_!R*qjo!amqxaYq?*1Fzb=`aH%Jakgz;gYL%%<$JLD8XOJ`uXM`PXK1 z=1;zB)Va>PR!6(`H^0_qXZoJwReb=u9Cwcym?M`3FdLk*8Pm4yYhCAa-M6SK$9~kc zJ`MbF|LE zVLohKGuorac+8i9j~b+N9Y2pdUf!nm1B8Cm_<6g20|@%bK=K{R00h0?XTYm}I1pph zqcD!77a)&qy84-odIlgozu5dvJ&WRlPicCyWSh;h13(zteeJ9{ZXaD^`o8)M5c1Ut-k+*r?pv~5Mxnz^~torrh0`P13>_LcFj`Ho+( z(cw{dUnxN5V;;}7kDzaP&vK#q1Q7aG<2ofj?VBoX z&zPm_onF_tgSyiGd(}38(9ZhJhO&8I+cjqOdd}0k0YW|X+XiKqHQKwzx~}V~3_y4f zx?H2Z-rcwP zX>6Huum}bgEPl^^Y*!lqLR%&S$!nbCbylB= zF2D2D=f)jh)&5r7bo~a_buQFBo4&8WXp(WS8@GAAKR`DzsCMMPaqvEj+n(Nsy=pH& za{j07#yxMm-g*~;F88{*o-fpu=YdiEZv7@y_q48e>u!gh#snaYP1hK``nF+XoqK*y z?+3`Rx@cFT<2~7juH&*%HUH8f^t;Z9h5k-C%uX2>)ctZ;*_`V$t=D=VpU>o) GP5!?KkmQ{J literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..643da261a16993d1f12c6c7e68aaaf68788de1ec GIT binary patch literal 2390007 zcmeF41(#jL)^?)=2oeacL4yW&jk^%RE$%`Lhv4q+?(XjH?(XjH-(20b?w#4=&_CdN zZ;eseHP5PD^3+qccb^Vn`yZ~~r*Wf3jpip0?Ax<*pYGds-l=_`PCIq%)qcQ$?p?Nd z`LFBunT6o=;X|iRoIER~DN_ee7&dtFuqKovHLbt==g<3Qi3Xi!>>Dj)npz4`YqtEb zTxYNFoE2_U;rqq)|EVfU;J+@cx;I`tGmt|8(w0~Ys4m$Sp9rw zzHds2wWx7ln-U9s8#S8wzCBBw4&xt+-+ws4JD_OVC0gpKWA<; z?`zvYY=s)JoZpHyV!6Losu9ckR<03CJy)p_i+vJl)w(}toyb!)`3&Gl-;QnU3Nh;2|KmUGy!MlAc=s75S( zwsDPEeOrwhZBipPsqV8zjW(?jYg1yI)rh6fHm?!O{@OPX>rf*$t?si%jXKtd?OtM? zYQ)lmog0XCsS(Tl*R@8hP2FdW8g;7?%Qfm=BbN8|s1X}g_gSMxJ!{1B8P=;tEcb2i z8nHHYpEYXKr-4}C8nN7){c6P8)P2^d(H1pgIrsiGV%g?^8nK+iz#6ff`=A=J^x2j* zV)bWuqeffRh~+b9>l(2(b)PkAv`vjzp7Yz*h-H7ct^WB|&+$B~f2P%QJi{VK ziu&=~s(*IXb3CgeM~ZknpCU&Jj`J(eu;7hK=291rGq(Pmt>-vb>(ADDjJgw(A zKkLuXdfpt494Xq5wXa|QdXDw3U-Nq28IBw&>W632H^FOO1KXk^`W#=ch~8$;mA#M2-~kIZ95Ri$df`5udx{->_3Xvm4e2J2im!uFmQpCOAgT3E_z2AeEEtyN*{od~p_kItKy&gGI@Cqd-d%s8A z`#sqEJ=ps_*!w-WUCCVP`tyE|`06ERJl;DaM~e8GB`2>%A#$XM$2(=@NWtrroE-0t zkt0Ps-WMZB3XXTg$dQ6KC^^~ZQN%YYIpaQ$BECt<8Q+vbNcM+)v!a&lh^kt0Ps-hm=V3f`jRcOHQ6fA#$XMPcJ!nR|=6MMSQoC zlXs^OIa0*;C^>mg3Xvm4+|U2uy-Vg&7vG0#650{+$2!+UzBK~N}$&XQp94X?Dmz?|rg~*X2{$$C?Pf>^*DdJC;ocs)h$dMxc zY{|*bQHUHV;?I|y`~ro@ks|(L$;mHKh#V>6FPEJB3Wdm#BK~U0$*)m}94X?jmz?|t zg~*X2{$|O^Z&8RGDdKOJocs=j$dMxcZpmEg_QHEyBS(t8^M1+6A5e%KDdHcNocs}m z$dMxcammS_P>38U;-8kB{27JFks|(i$;n?(h#V>6UzVKw6@|!=#*@D;IXLQnQ*!dR z6e33&Oa89p;5b{pFFE-K3Xvm4{Kt}mN5>z$2Y0T!=68kla1)nh%ja>qk-VJv2FlBe zuczErd>!SL#Me?@S$qxUcH*lkuO_~V@+RUdDYqA2LAitYa>^aWmr?E{zLav#Gd{kA z+$9PnUCgy>0@sVk-NYAC?k>K7au4zOlzWQLqufh;F6G|hb13%_pG~>1_$yGLrjsdbmB95Ray-j3ok(Gu1gC9wI)P@=)2T+)hz;%D}MDc!u$dMvGPsz#iQivQW z;;lZwaxVnUZ!L&b@64nMvfHm^@(L6pM~e80 zB`2>$A#$XMuUvBSDik6|iukG}C%2;zIa0(|D>->}3Xvm4e2tQm*Q5|RQpDFPIeBdg zkt0QXosyH+r4Tt%#Mdi1d3_3zBSn0Jl9M;25IIuBH!3-KV+xTYMSPQzlQ*RhIa0(o zD>->{3Xvm4ynV^Z9VkSO6!DHFCwHO{Ia0(smz>;%LgYvh?^<$lHwuv>MLgbbB1a1D zQF3xm3Xvm4yjRJ|y(vVF6!AVKC-kmYlp5g~*X2zIDmT+faxcDdO9foV*={$dMwxeaXo?P>38U;yad{yc31Uks`ix z$z1AqFxSYDB0i+#6BT7ylNg;Bih-XfY_oK*>A{OsQks}3< zEjf7{g~*X29?#Iok%A|boIH_2pRC4kz6e358_~eq4r%;F-DdJO0PM$^~a-@h) zFFAQv3Xvm4e7BO5cc&0JQpEQtIeAYCkt0QXuacAZrVu$&#P=yVd0z^VBSrlFl7mBw z2TBeOogOSX`5_9CBSrk-l9L~y5IIuBA1yigF$$3*Mf~xSlb@gvIa0)*EIIip3Xvm4 z{OOXDpP>*rQpBGvIr%vXkt0R?`I3`gpb$Ay#9u5q`6UXGBSrk>l9OMd5IIuBUoAQL zH42d6N0gj=B!$S4B7RiK$wyO&94X?*l$;#j+iV^{no)9; zqwfDTEskZ*kt4119Qp9g7Ir(-9 zkt0R?j*^q_q!2k$#P2FO`ECl4BSrk4l9S^sSUrFg`yg_psQbUYP`@8ikN@q3|Nr;G z=CRnM86`(K-pBl04377M$dQ8mju3oc$z1B$ zA=vK-!G1>w_B%qb-w}fSju7m3gkZlT1p6H!*zX9zen$xQJ3??g7dETA{_iaKzkU|{ zZ}0#2Js|eR|AoEZr0o58u4j6PLOh2vJy=(mv2Tm=K|BL8Mcp_vGd)2e&csZQQ;0J! z(_<9kjLY;Wg*f9fJwhSQuuKov6=weXM^j7i$r`cACH7Q}*z^*6x<+hVi9J(d%?b}I z*SNeu9w~mFa{U=v|KW37Ml1eoUHQNNPQg&c-lpuf--1Ud{w8Jj^#(jf@z*KmnB(Ku z$m63>(yLr2ByfF&JW2dA<^o=qU~mSw{!I2d`~+L) zA7QWI53qIr9=7h^!5-_kuyy|iw(eiUp37J8G;*ddDeRWO^$W7+_IX`w#=i5EI@Ue@ zDH6}EIqY?41`jXy#ic3ib!h@yi^j0Eh`(Uxx&K2Op8Mah=l&Nwk(}vI3fAHec$W-v zZ9)6?1GaXHz}BuMZ0%-;y=P{F^E`}c&kDzW zUmz|dYqvVFe)S*M{ka<4U)&DfQoJfWSiA~sU#<*W%avejxgu;It^ixhwy?Ea9=4Xt z!PatF*jg?FTg#&bPS2VKEIuf&19bjwK z9`+vI9JY3w!Pah5*jjA@TbqqxYqJq-Z8n7Mvre$J*Z{T`>%-P!J=o{ly0Eob2i{d) zSsV8Gw-#*e)`Y$1*MLLIJP*C62NAPR2g25N0PH>6AGXi7fUR9W*xL1ltz943+VzI* zqh7GJ>j_)C9wdo9pHm%}9YFB@$`*S$44&q^OC-G2tpm+#8 zSUecEZ+C{R?M|?LyCZC^cYv+!_OP|x4z`xt!q#pZ*jjB3Tbr$5YqKRBTC^;4@Yp92 z^B$TA`#hNddu-!julqRIV;u|oJQ)Mq2cu!H`zY9R83}tXBj6bOY&FKdJ29UJyTP8z zuCR5V4ts9XVCy;+_E@IC)^RfI{kRJpW1FLl&7L}lm_2nM?EQEE>~+{5wuknE?V%a4 z&$jr-&%7@2Z@YOed&AawFW7V26ZTy8fMYJRm$~>EcnmR*?P%CK9tC?WN5b~V5wJaQ zIP9?=279cB!XE1(usv}w9Aj-=W31;Bv+n1?*7;o6x}F2uLubSG&{?o`KNI%eJp=aK zPlrAC(_rt9Q(yrfdtHx%eJ&mghjz8jf~$ymov(zg z)fKR{xg7TXx(v22E`>e!OJL9aV%T%P2=?4Bg#A3c0FJRQ5f_p@`UJ6d^&i*$`8e$T z_!!(#@kim_;zwXVKOctsDgF?=h4?{ufcOD;p!j}xYw>;XVDY`MeR&UT?eB)I{avuN zzZ14E?|{7)x5M`7ZLocMD{P^_8IUBY;9kLt>sIwwR{n_mM_58 z>Ur2&JqKH>XJKpg4D9pnY1rC51zWo(Vee_r?Nef&+b6Jn@-ghWd<1(gAHw#<2e9Y% zKJ2-@2YW8>!XEoOu-Ex*IL6*8E+nt(ufzt|e_Z$HFR<6~XV`ns*BKgH*S2aD%`t!+!#KAIi2k7k3d{j9LHZvk8TSzzy>=CGd?&0uTa z6t?f0z&?)}!(N;C4<7xqT5H&A^AFm3ZT^P+9Q_OSp8FH_TKxfgt$v67T>TA>wVS_a z+E&`Gz*stn+rmS{%fr@wIoMh+3tQ`DU~9WHY+o-0TkAHkeZC}YZI^(p?c%VtT@1GG z7lp0uBCxex7`B!R!QKZ8!q#>H*xJqyTif~I&~~AsnYG=Fn4b-s!h@ylCh!pP#;~>C z2)55RgsuGsu(e+ww$Imtt^K;NeZ3BB?bn9w^R-~F#hS3SUjw$UH;26ztHWN4)nKnh zJJ>#7753V!0()&%hP@Um!Cs3M;aH3Lil)|jAh99RdH`&#`@`0H3)ousgROO6*jo32 z?fc%aweJO6>z=T+?g3lt?y!B|4Yu}OVQbw5w$`0tYuyR9mK|Yh*#Wke?O|)#PTGwk zW^G5pLlhqYTifBVwH*dq+o7@-Im$i7dl2s<-W~2I-VOH8B)h`>6`u|d5Kn^#il@R` ziKoEZh$qAL)Gn|+F$wmXPlUb36JW3Lc-U(^4vsZjy!5%e`aYAGwLb&yp!n&qwLcBE z_NT(u{uJ2SpA1|3li==Z|3tWl_ypK%b3EKj@#Em$;$va2-7#=q#gB%)mPf%}%Ohd0 zJs8!tBYZ;)kUz^>O$CSbpgDk=6^oywL1^?K0FuRR`t(;hltOHM~KgYL;GdRXWmyY z6Z2ZU1h?0{FTx$fFTh@#=V7nSb8u(Xe-`#yJp+5+Jq>&9o`Svao`k)YPrzQw$6>GK zW3boqQP^ww2<)|d81`B|1bZzXguRvzz+TJyVXx(V@F30qUf64U5A3zQ8{S&=?}EM7 zcfwxlJ7B*D+zyY@{kOqBlW&E+uWy0<9KIR$KE4Ug=VUZ{BfP8HzX1*}wJo37OMeo} z&m;T?Y;XMz+grcE9ZEZK`4zUeeu3?+pJ99JC)nQl5w^E}fbFgCVSDR4*xvdUwzs~4 z?X9n2d+RIM-ue=@x4wYwtoeHi`V_XeK7sA6k70Z3BiP>h5Vp5IfbFgKVSDR6 z*j{=U9;Wqp2lhK-)5bG<>uuup)?2W>^(Jg@y#f0f@jC2h#A~p<_bR-*#`6js-dnAF zW^b+3bY^d@2zSuEE5Pv;GyDH zu)Q@mY;Vm4k5K(N;gRAwV4rg>VV`re!(&u`Hh8RfR@h!`0gqRF7I=cVIczUCgC{B8 z6t=gU!1i`y*xrtRY|GyM2QN=odw;_|Xa9m9P<&; zFRj8k>%FoOaqpE4Vef?vVDE+XVS9f)*xp|kwzt=T?d7#$dwDI`US1Qnm)C&p<<((( zc{O-9`Kul5bA46VKfA92hu7CA_t@))5wo`sh3)M_V0-&uxTE?$2)4HmgzfDEV0(Li z*xud`wzp@%_V&K8y}b`?Z|@D;+k3(G_MWi4y$5V>?+)AByTSJMuCTp59k#cp!S?o4 z*xsH3+uM_2dwUny-kt>8+Y@1Xdjf24kB5h7J;uR9#baT6e++EzkA_F6{wUbbiIK4P z#t7KYh2gNz{bBF~wKo*@-WUS=+#d{kuj~wauj~YSuj~kWuj~MOuWS!{uWScc^B&mVyc_N)Z{7vln|H$Y<{hxTcROsa z-3Hqmx5DGwG7ef1dZef22pef0?Jeey8uGyWmi`{qH| z`{n`IXZrnccxACNCZFMd5wkb`guS+Zz+T(mVXy6Pu-Eoi*lYU>?6v(F_S*ggdu@M& zea8L(d#%5Ry|&-MUfXYBukAOm*Y<1JYx@=Kwfz$IKKlY5sI&gT2<9!d~l5;7;m$W7uoE5$v_y5cb+`0DCRhhrO2T!CuRCVXx&n zu-9^J*yq4nu-A4?*lW86?6qAT_S&undu`jnJ{MMnz1FM1Uh9=%uk}i>*Lp>Gu=HC2 z_BqlP9;W#6u)VPyY;P}SFfaGWDcmOkvAoru{x zJHqzP4zRtmJ={rS-46CSu`O&bZ3Ej&Tf_F!RJ8gVy-b*cG-Hro;BaG}vC43fl`)V0&RQY%lBr`&^j>+Zz*M zdt(CZpW(*CK4-?kK4-?l!=-oV>vLobaeHYrY%h(1?WK{h&y^9d&z0e@y)+EAmxjWA z7aamm*Z2m*K6iG8!+T4WI`-ZL#O%HE;r43lJlI}57q%DAfjg=G*|5EL7HsdG3EO*T z!1mhdu)TH~Y_FXP+iR!5_S(s?y>=38ubl|nYbU_=+VQZxb{uT49ShrA$H4a1(XhRB z0ba4Ujv{Vv9SIMXen-H=#D~LU#fQNY#D~K6PW8X1IE47Fs(&zSZyW@NH`Xk5?2Q+R z*&8pwtEsK$VSD8{*j{-SwpX5kJL&$XVV_%1!S>FRu)XsHZ0|e{``mgAwwE4-?WIRx zd+A}=UU~?&mmY-er3YYp>3-N=ivK>uUb>ICy>u^ZFWm$C+`1dKx9)=Ntvg|R>kio7 zx*fKcZi9zN?_1%a;#=Th;+x?S;+tTfb2q|1=Wc-gGxhbb&%Ntld-qz{=ioK)c=dNR zY;Ru$PgML$*yrXIu)Tjd>~r%n*yrY@u+Pm)V4s^8!`>?w!QLwu!m(GDEo1Ut`I(sa z%1^NO#*eV~#t*Rf#`kb1&GkFjd*NHyd*K_{d*N%?d*LhCd*Msid*KV%d*O4~d*L(K zd*M^q-v0!)_dka1{f}UK|3lc`{{Xhv--qq>_h6rs@51)(JMb{+^)~EhboIaId5gHc z`X+2|z5&~tufz7{Yp}ieDr~R40^4gZ!}iuou)P)kJx_RPr7{P5sZH~ly|g53FD(Jv zON+zy(qeEYX|pKYS-c2rZ!HYlTMNPV)`GCTwE%4I%n#c;^TGB`YuM-Lys*7A4{R^B zg6*ZbVS8yV*j}0w_BlESY;U!M?XB5ipQE$EK1XMT?Y$PT&(&FApR3K`VbZr5JY3uq z9wBZ5`y6cy+nbGGd+{H~nB^*j`=_ zwwKq1?d5e~dwFfxUS12fm)C^t*xp(O_BpsTY%eVZhnLnaefr$nm6*LW9k!RI!S>Qr z*j}0f+e?$-&YIgUu)Q=1wwETt_R<8{UK$VEOXFaBX)J6nje+f@(XhQV3bvOdub?aFAafxE)It6t({?eYbV&=+7Y(5c7W}z?O}UsJJ{#qwy>Ys+rakT z*6?tx&sMO{$t~fLiVuQ)ZVrU))d8@d>HT4!n_IvW)LuW>=V)Kp=V%{z7u9dwWM=R8 zCT{Qdg6;jDu-`X(z&>}o!#;Pr!MmxyuCVt`7kCfFJHxS;RxY1;FP%@Ut?Hf!dvBc! zdvBcsdvBc$dvBcudvBcycUJpnz}{P@!`@q`!QNY^!roh_z}{OY!`@pb!QNXZ!roga zz}{QO!`@rR!QNZP!roiQz}{O&!`@p*!QNX(!rn_qz}`!T!`@4W!QM-U!rn`Vz}`y- z!`>?g!NaBZfw1?=0kHSR{;>DPe()IGKLhq&*cbL**a!AGzc*~J?*-e-d&2he9EpCxARZpOZ|cb_3{?>-INyHCOP?vt>+`vlxsT09QhyN|*4?xV21 z`v`3BJ`CHt55e~CgRs5(0BrBx58J!({}l|4q!^{|nA~ zao3-)z5fSn@Ba?l`@g~V{;zQ7XsD!LV0-^(*xvsMw)cO8?foBMd;fdb-v17^_rHbh z{cm7<|7+OZ{|dJEzl81mFJOEBbJ*Vh47T?_h3)-MV0-^#*xvsLw)a1T?fnm6d;fjd z-hU6a_uqx>{dZt{|8019E)LhX;F01tVSE1#*xr8~w)bCy?d?}#d;1mG-hLUjw_k$o z+6&`*z0S~GPBp$fbI3wVS9Zw*j{f3+v}^s_WCMtXKA-G z+(o<+Y_G2f+v_X9_Ig{`USA%z*O!Cs^<`mueHqwZUmCX8mxAr}Hn6?EBy6uQ0o&_~ z!}j`Ou)V%0Y_Bf@+v^L%_WDAwy}lr9uP*@G>+{3*`h2jx-Ws;o=Y@xBJ?4Q&h+Dz- z{@k#=KNoE8&k6e-XAXF*+G`1WZ_Ez+xiA|%LG@>ay?0u`lN6r?-bK6??7h_-_TFj+ zdv7&`{oH5*`<<#W?7bKNgLdz&f6Cr!tTCjXYg0G&ontwcbxN#Bxz8S7kC;8PZjE+Q zgY|2~(hnPySkuyneXubxuiHj&jxDNh*g(6R)QIJHHmwm$zin0{mR{PtMr_YA?)D|t ztZ-a%E**$1tIs>a_lrBh?~6N^w!9{7O1w*rSZdU@Ml8qItwt<8+r36C*QZB~SbC>t zjaZJeR|B!$4aEA?h^6=Y)`(?){ThgEQ6rXf?_VR99ve_2mivBSjaZ%`gKETbjkc^2 zOFg%$5ldfgT_cujv`vjz_P1>VvF&QaaxB|75Zj>=^I5uMB^I{LvuLLVVmnu2Ik&if za09U+4a9~v5F1uvdG16%`F=8zSn3|1kAPE`;Ng{a|cp_eGjO_tnY!9nAiTG60;Xmi`WZp z=MXs0o0YG9zo3SXh*{N+=&W@!#RfF!(eN4C|unGxlTuw znAag=vA3+*(WP#VF~rZj*nT?}&T}blIR=g~)?K1V%C=AAL+g4f z?6yyVea@c@TlbS-_kAK9edU~e)|^4i+MZtOHjZ0J_Scz}n8$QhB^LMQe5&WL?>oEP zS8HtN67xLIsr2EW1J0|olkb-?jthu+UC*!7^;j<=W-nb>>Zbp_E|(IsmY2XD)5Yc9 z)YAQ=|7*oAYv8`iD>1M86(!a@IwE@J zen#H_yY1^?d+fT>R?dAyG+ok7C6<1$p0^P5d~dF_3Q5uEc=i4@2d3ay?jrl9gip8G2G6*m3F))_f=vZ>;08j z|LCNo2P!eo`N2vo-@)V7hbpo7EXP{)w7vCkxi38(_vKh0A?CVIz}D(<*lYP1?D0Md z`&_hto+9Qm@|9`+*?bF*w^mo?NT?j@R|QkC6>>znAf`{R%@<16SIdtfJeklCA|;3 zt@q&YUydOtb*sJ)Tkj9ceYNHp&lRuhN0mNCL?JIP zX*<`_vF}SP=RHDt{YcF3SU;4y_J58seui+}pJBfT|J1;}ajtl;{0isKUNM|s%Dp*$ z|BUrpiRJjMY3y&GF~3*Z8LYYfQHfdKKT9n4)ezm+sKw0hRsU4_{2R`F6Sw_UX)8WU z4P#$)Hg)x^Z7)*v|<2xp}2-vr1jhuLUvtdzMPw=q3A&J#0;9Ep=;oBHk5! zp3YwBXSN3ViRZ5U)Uwn~55$;q9I-z`ha7`7oQwOsUUS0MX^u*}InQWu?n=xaZB=6F ziKw6B*@f}+iw{a_O@K1)xA=d8~Pl|EfJd=Z*Ak4s6}PW63i722_PSB5>tmEe)t2-g*1uctkrI$Fa?cS@X?G-KJ%(9&1NpKEFD^Ui2Wd3+#87 z&gI^0*Z%KTiFt0_D=~YmMILd#@79we?=;L(JpuU8x&#G_%Y<+hsZD+rJt`04+ z)Ys!3UTJ4or5!)>MpR;MXJjRo=TA&&R3#Rl;L_@3h5?#p+VW>JM?&rDz})?+-J&!MP04)#0C*vi>wbF+*m|lsw|5)4Vy64zp?cC>X zl|FYZ?dJY;yRo*`YY*z$GrL#Di5%}vH9V2A{H`rM5kt*btdG~=DEi2?i+V>^#^-yh@4`nnFur)#t$HN;okYz0 z|CmZU_Uy4Gmgi-(n(f5*Us;aNPoS=~J|6a&cwA{Kb@9C?6SEE{!5+tn<=(n!(O1Xl zu%t7IdCkv&y{@Oj{(1Z~*yr!5u*Z1{9NOhLgL15KHhC|cO-}(Ushjif7K1A3Vq%{EMexuJa=j3?9v8sY_I%jq(s^*qKi_fO))keuE{E-< z%V6tsDePzHC9r*&+Qg?<5wn(8R_aEbtRLsD&#S9T-CVP*7x!F8EZ@-B*c4C^d`k7wsLl`2hC~ zh;T{wSK4vx!3JUvm00fIKG`7GM~GRwhb!$2RNY4_G4GMbDlu#NcqQh2`a~t>u{>E~ zId1#%X=3jGsZ!TxdY+H*PUm}{qi*US(|Z>7n4W2%t#}7?TQ64HdI9#i^L%M5=iqbX zWn%XHOO?7Ywwyzpv+nnmQa63-y4CkSpFgiw?(^Pwt%3e`uJr$Ur5*on@QnuAi8CVS z6k~p~(oX*T6|uJ}vG(y{NpDwTzVDq%tbg9b_1#L$ZN68DS-1B~EY~rb99i0pcLw|Y zgGxKDn?GOFRVnvpJMkR$dVEGDulc92J^u;pem{o2e?Ee}E+4{i)-*2t#${wFhrfId z4Jn_;a^E*lH~i^#e`ui3A1krwYef7(s{YL2Hh-$zXKjBjv8MIS)cu)l#@UnQ z2>eQ2>+?&wH?@p9SwFnxv*x!3>Q>+BelK-Xcek5oO|9|9&zN4%KWWGQ{G&1k`{r+A z`SV5e_g7^Me)svO67xHEqgmtbncsU;llYm&`>P3c-EQM@Z!JCJ{l$BvX{pP{|4G@; zI*iYAo0YoOr3LJ{&H~#n&0()yGuZpM|9^UH)M&O!EJAr^RKHvC`8IpGFZZ-{nTwd; zzvhHTsn0oJkFO>vC`NGH&a? zFfsSD5bX0}kxE+&!q#p9*z*la&0{aQuf?eAcYsC9z3DM~cL`!1*W#6aMn*>^Ek(?_ zwt=n3lI7mi)&5$##BvNC@3NJ4mZ`K8-RE5Md!y_l{U3FftMuVmTVg&Jmao)xAJsj( zLaCeQk2P4CnDtu;_PVWD?oH2oKd(k?RCHWYJJ@Zn3R~Y*;P3#;{hySϚ}v^I4! z7M0h6y_RdjK8MzT?d8>B?|+YDol2~4R4r-UN~~iBxUN@;bx~~nN^F2)8&qOGi#9B= z^i><(w=pr#d80}@Q8hKH{)`l{Jdff`vtKtU?W7-EH@+jXA2y?|_ra#+-W;#*jpw<~ z@XafA-A)H$_F4N#_|o@58N2-DYu1Nja8yFY=mi zPhHP{JJ@sDwzO4i{&D8odpnf6c?NsEcPz15V~X#7J)gmqG3*SxpPefG_&eqym43X> zM-cN|hr`~J!(eMZw6vA;2`N*@I9GiRI+ouR=Y7_DWN9=rkcG@6*t%u%JT&y)0S zoIP2N+s9JZdXFjh<`}$R#+6u(!9JWo%vz7H)Q!Hf-|F9^m{{r4btjdW*FW2dpG&>| zQz~Pa4EsH8mj>FZ{+?_)b?t#^u-ANQX{%`rwxk?a{M_O_yBl?_)vo2<)G=<)d#h)| z?hVw9-=kUcJxbkLTE)*i>A~o4Z|Zt&_kw*E?pfNZHNQ9ua=g*jzLmE2Dfgy_yw7Ho zSe|upf3~$BG4Hd3VC!}u?DOgX*k|AVu=kedcW@=VX)WW zP}pNUr1V>BP2%VFoO_JtDC)YeBP(qkQQFFJ`g}N^nAi6>*xDTnTdQMWKkJT${hSEO zvBkNR<@o%BN?q^4lZbi#CziVDg+9?hNheogZs(LrEN;m$=J7lRZ1FCu5nITlXIzS?a!&S>vcS@67zHB z{1Wpyl6}TGV*g)IX~#c@URa4`8!`NgN-XE)v0X~c`}GpodR^SWy;l(PdS6~?>#}lh zuCeQ0S&6yNt12<;eszh}(kspYKL@WVb#vY$%38+ddSagUb+Glkwt;@Czi+ys)XjND z`>B7vE7Tfqbx+<{>9cJ#RMO4Ftlv$QK7CHyQervp!PyAc+lhI<+y?vkdTY72?m~Uv zdp{ozRa^JLqr~^Z9>+bfpSO3zK4E~V8&x&_o>-skAIlKk?*8R^GZAR)0dT4%qeyIszxl&)vqft z&+(f|ta~(C(zlhEwfwFU^UqA*S7JWTf2hRV=8u(FmmCP!pDM9!6#Kan>#W!>C02Lg z^Nx!B!-uZ>H|#$Dg1f5zpRo1$1Ma2x@33|F4IZHQudw&O$I^JVng7kTMzB2|l{(ou-DA`v?kVH zb>}U08Pb1Jjx*jrJjaEoYi}(Gdz=fv1OL+s|GgD#&(8UHaen%JqkBkZ>t<1fn6S%Ha>Br~&DwWu9-M4Bb z=4W2J63g}Wd&wHa+|KIoXtldqxwmoo%-&eD#A?kue%9)scGjV;eY7^brQ&Oqwz%y- zDfNx-gzSqAscTVlBDPwW=dK_J;~jp9!DqG&yS9<_eKZU<7^LGqs`%e8+ZKO z%ZR0af^yv9Z|l>&(r&lXZffm&l~ex6#K(% z_bK=0^Czn1oZ{ay*d}k}+P8ta@g2Iy)UPro_c?%=*RVhAHQAzaug4VMvHNTtL|vbC z1IxX&bdG0aTg_)1>UwW(4O^?NVDGOjOWU<{j%VodYI{5CdQG=2_tqL${0#0n?L=Mc zvmD zLVe^|;+g7sj-#&o9t(S3V_=_WqhUYCM#1iTBpiLWDt)&p*SPFO%zf_(d+qmtt<&zX zb=nQKH+F@+M$_T$>TeqCHJb{14^M$T*U7MT-v#!$H3{}}VL`0Vjv zN!d?4$J^$ITu-MRYkOLy-&0}F=@i%=IT^P9PlE0L6Jb9GPJoB2|Ks7%q}Dv**&VUE zDp4fmJmT5xd7MWj`|VuV-%9 zc`0nIE`hDp#c=3RtM7QmdL6H$uKT_g_RnC~z&^jPhP`fA!5+(%@V|K|p27CaZPfMJ z+zNZ1x4?cD+zh+#n_zqYM%Zg`1H5$%r=;s)`^0_SNz85EQEBV;a&Pm}R-1B-%YDSU zi0_5n_C2uYaW}k$>fZ&&xN=W=zusSB%}RZ*-NVF&shx-5QQ`;T{LUagexS5ftDi?I zvB_%Z(MrtYd#uEAe12v=MQnuHc@p;bcTd267Cc_s%JJpEzDZC>$r);zCIX~)mcPhg)TAH&|; zAHm(D`I0__J?{@-Kj+?uy+_`Iz1QD`{hWITj+`><>$&^|doF*% zp35Jw=kh!3HT@0tn*Iv=4E+V}uKs_9J@=nrulbL#&yOEqYx6zq^W;0&&(?2Yd*mC~ z9{C!!mS4esetZf09Q^{e)}O)}krwxi^75_r|d2-U#+t@=uvti!w(0Xddd?H?3f=vG8gQz%n5re zbHLWUB^ zgRT9lu(e+Ww)QK--a{+FUW*lBuf+RS6=u(j_ATl*fc zweAjE>u#{M?h4!2U0`e18MbzvU~AVA_H(NPY^~bE)@pOuT5Seft4(41b`v~$OhdmRVEp6kxA=eiT@x$X#it~-MnMbvxL5Vq4f+Yy*3pw}!3FR_W_1OoBbviExkl+v@(D0DErZVb5h8?757EJ@zrMJvSQm*hj&Bha3sV z*yq<6(u46Gg6DD&b-m69!k*g!u-AEi*k|Q_u-APCY%TVMz0Ui`fUWiMu(du8w${hOK3k80t^LukwLS{A z)7->Y#-hR+jqCZ*5(%2+T09Vi<@9;aU*OkZh-v^z8odf>?(uz4kNRqLe?Cpz9)1dT-%rB!Tu|=E z`2N*>y+vJn?oHS_zX5wJufz7>YjBU~prlt}ufr>_*WqQ@9()P*Twa8|4llr7hv#9h z!*j58eir^W&&BU^?YTo>d+tl_wKiYC*5Y&6T6_lk?D`b8HlM)WpC7}X`$w=n_#td9 zK7g&o`>?fm54H#2g?%=?1KV?N!=X*sQy0sZ|YjBzhG$@kQgrjZ zjatsE^FO8Te5I~+oR7NJtu<`j=7sICd0^|<3bu}O!`5*wxJT(XE_1@xbq?6NwuG(g z?67s64Ymhog{^A~*t*UF+mp>b4*Jd@?Ytat2_N&5EG@n)ASc|#zS?-JYw@dG4Kw+n1eS`>-P%b6v2^wN1IkWjHa%4gFR+l|;B;$16s2P?J*F^^;SO5LHVyBD!B;yp{<)Yt7^S|>v7@$JKXyQ`hO zEA2!m=aauHZL0bii{F*n4>M@TUatSrZ{~mBCGU;jRkl@|2XU{xbYP{e17Poq{b8@q zesHYEtZFm+ir>xI_eWCKXUGw-&w#^WkL57fzCIN8SPp@`M-PVWd!L<0mDrp$###Mc z*J;%CS$Hb!d7J`!9w);-3r~W*J}1JS*9oxK>v-7nIu7>OkA=6?xQ~JD-J{`{ck5E8 zz3N_0%yYR6_FOK7J(o*h&*eDQu#4Ken7GgNi(s$eg|O#(0qk@6eAsJxEbaSTK99Jy zI2ZOBp95Q)vtes<7VLBROxW)SXTW|xI2{h{YWX65H{`u?Gj*-)O|Z4Q5w>#x_B`)~?UB1+d*n{oI@|%#ZP1x)72JE@K4tp-I!7-QSaVaVHQv7{U>-q_G zt?S3IJ@gUmzCJ8%H_L{(X4~=Iu=m;L)b+SNgQw`ePfJ^QK3KyqiFy6LD0OGk{n>AP zH|q8Hk-Bc{2iWWJJ#78IgFWtVVfXzFZ2iB6t;tvLzj>$n-Q_>jwcdZjJvFDlVEgz_ z*kkzvcHh6l*7!Hrd-PX$OSS(C?DOYmc%tGz!7-P)^jYd0-!=Mq)RK+weK9+1-DiWX z`>e3%-U9BWerAC^_vWy5KcD_R_h!UB_olGt-URluxH0UxH-clX>A#?yYyA5=eop8A zx4TZBs7T7X)&H(%KH9YotzrK>H80$&zKOa&=Yc)%`2T!+&07(-4s*lSVJ_J7o)h+( z&jH6g=T)CM_Ud3Ul#TnE(7~4TN?JWcPZFw+y?gCmxO)3ECGA2i^D$q7K1(a zMPbi<5!fDB7`7G*!CvzPVgH=G06ayy%nyfFE%aGxRsF7PW$OCeSqZi_E5g=d1=#D{ z7WUkhhokRWXKQ>nYE!ZAf6SjxefUWE5usyOGY^~bCp?l*p z2G3>l5^Gvw)~zEk_tOD(KkX~+4lZrPrBkI{$GQ?5uDV@dpV6Jmz3G!ts@s#;IB^fy z{_YM>RlFMAUu1yg4Bkp%Wkmy+7<#-pjXJ=5?dd3e{9?L$kJ-9dQ zdEHKX9?M?DJ(fM;7)z7Vzdg1;F|W^lr7n-1x@3(GCFXPP5ZLoL81{M{1dlB3#N|NP zYkvph_VGF2$HPk9^hw;A@p$+4xQ?c-=X?}wpBxEW&m&6PjiXVLwLhj3 z^Bz35#A=Nx-W5Hj6R7JoJ|6a%j)QYPF@}4&-`)+%c)ZKG?NhneUOlDK_Q{pDPlCM% zPlThdW^p08?K6n^d_5huwx_|K*L`K&sb$;xChPt@v-C5UKFc{(zc0Rwy6si(Qn-Wo z61b!IVz`U=BG}LO3*l~xUjTO(pAY*?JP-D}#<{TH0nUNFSI&lghMon7rn#?za<1`x zwdZmTbv>7>VbA3%*kipC_P)FV9uysvbUAGASjU@)xvv{xzoXp%d*0W>?&~_(=f|~h z^p$Js^}dgo`@Xg^?|X@REcd{jqq&mqhONV0u-`B5guP~Wzn*U~A8&?ZuJe~ZJlB_rd9E+PUgH;G&-DekQ)wqI&%>VkbFk-rLFXo@BNjS=lx5it)F3g=qK17`VqE=et_+v?_r;b-@*3Ow{WkRTuI-+ z_R!a`b^Z#r&R@cw>ld(f{~WgNpTYLfr?A)I6WBlFe+*B^!EyZv_WQ+$uss!&>lEM5 zd;d0`bLMqw1Y6sGXv$O8OJFR)4_O;&<3u{07HZ=Pmc!2dx>4eK2oj ztn!8BcpaA|X04Wiz5kYmy{=2a)~*d~?UsbC-4d{$my5&JZZWvG z=CUYk?G}Np-NLZ7TL`u`3&PfB0oYp15Bq&=J~*_QRbxoc$L}|+)ymYhHY>r_W<}WN z$qKNwY72W^mxsM>%fT_0x=R^PjwOC3@>tfTl06?kBYAA=5ck=%Hr!itTno0&Yr@ub z4cMMs9k$M^!PdDQ?76H8TlZDqm|Js=A?FsqAMso^q^{?(0c>5@hppp!rR};)X+PVJ zf5YCob*7U0>IA#5jTQ9h`#?lk^-su5*EZt$RaW~jw>k4~pU0{zTD0Qy>J|}*rw5|iV*XKZgxOX&A z(iX7C)(`er`obPdA2`O6=YLR+rTY7rVbt}S4~0FpA+U8F4BKxz!`68x*n4Y7*w2R@ zV0&bHxOWV&r0rm@%eJuRz76bk*&4PMTfx?1OL$N=&UFy%y*Cj4w{?r(=UAJu)b;aV z3~a4N!`5OH?DKylY;8us9_w&8#xhr2O3HPs{=R58>bmb;VfQ^9cHh%rKYOOa9?KNi z9-It&ukHeS9VWqEhl#M)VFK*Ajfed^lyUICd9?cbrNgLe-4BI5_d{TN^kCRp90Xg7 z17Xko0N9?~AMTwCz;!>^o}2+&i+Ctmi+x~gu{Ufj_JZxnJz<|~d%*VS?r><;N@Gf2 z#P64UhMz=Tdp5pvv34gCw{|DM*6w)N`#Zjqu~zY&OmB_jIOx+??PWfUC;eu*mJ)K_S`RoJ@*S>`{;bwd+0pa+MEm9 zN9Vx3HO{kP@2Rt3Yj-AW?aqL$-RZEkI}Ns0r^5ElDex5ae==+zwNdOAV)o6=u(i7h zwstqd*5U@(=lu1s&y(w5YjZ8!Tm4@HTbrw4YjG9q-|=1vTZ=1T@3G5aYjYWFZ7zjF zi^gRPmuW1wmRR%h*tN_~&% zK4PQQ&b1FEr z-r%U)G%h6Xk7sDddOcm~_o>oPYU?$5yV9=ve2)99)w|U7_Y}{Tc60pRBQFqhKhIaj z;PJj#iFL1^V%?uFRr>J0eudafnpNsrqt}SJ-B(N9+0|EiD*pbs?|r{=?|ZP%pLb!O zKkvZ3)$ZG{&zrYkpEqy9K2P3&J+If{So>OIt^V6upHtUk{|xS;v3v?!huFU!`zNr+ z_A%_|!bh;j`XTJGegMZ<=U1PpYyAGe?+<@bx4r8933m|x0eddL!#<~egZ&Ks74}?z zf&DJ_Gwiwl1lvzP!oAi153twmd)Uv=IM1xrcf_sLx3KrbH?Xz)8ut6fSMX@{|0V2a z>KAZmn`c!}YCDOaudQX%xn|a~32f~e!+xeVg6)HU=*wgOyY$^EIxQ*tuKxD{^Hln7 z1$!OmhCP~(7id#tm=UdP#BugjCHlh%IwW-8Y8q*^OZ9z9H>(mdnPJLnP)CcaParB0*TQAtU^@RO%P7l~Rc89HFH`rQq zg+r%W&-(c9E81!-BdKeTjeza3;jp!No%_AFhY`2uhQiiv2yE>J!`5zR*yrR`1?w^rj|Yc&?OR%2lA*U|74 z^*0I*t!lj=SASo#4|VOsyY(E_f+fT>9KF5xR?Wd#QKGNYx*yq#{u;+a^Y|RgYeV!i*TbD!NnA;rc zBiEq%`;#-N>$#o*+aDiP*6e-Qb3dK>_Rnds=YA^ext{`ijZcO>x07J+g%jbJORaMv ze*fY%zlyq^%ayP_a0TqKUk-cSE`vRnOJTo7sH<0MQ|T!ej)6=bOG$SpAUPT z&x5VSxv=-vIk2@k8@3P5fwkMx~t<}@8wR#G+R!_p#>Iv9dJq}x|$6#yqC~U1BfvwfU zu(f&!_Ot6j*q(dsGhpp{%=|v}ajCm- z8DE=njmtm8yr2JueK!0BTjxJv>--07oqvZtmfv7|_*dAv{{mb0pJ6}ue}cUZKf+#z zA7HP;_i!K0?>pEY{}%STd;|OW`Za9Ne+7HpzJ$GQU%>YK=djoDGuZ3+DeQIp1derV zT*f?2V{bHfeXMEu%^Hl17%N~g_H7k8h zRG;4y^O)x2UeB*J?D@?LdkrV6t@%qV=N_SKJAQxTb{DL)yFj_Od1=q{S&Y~a^|L5E zO1ua>LA)^Rz7~SLCf*ZE60@#Lz!TNh;^p34lblafS&G63elUR%``gQ^jp7?O5*>ODz2|M)$2s%=={(*qW>idoC-L zwo(&o-LAxPyk4`_EA6aS>ZZ3m-nEEL)EL)<-PRiA-t4EX>aIh~dvNViH+>jYbACs$ zSK8~|4XEqq!}_q-d_CAd>#tkdPAyz_!xF38s-=(YxB7dk&8XW)J59&z2MFn-G!J*y4MKgQ! zC}P&?NZ48(0b8rXVejX|U~6|MZ0!z#y`K+;t>ruvUGQ&uEnY-3fN-MA<`|< zOZU>o*F?nb?hdd!v9JTXySrO|&)K~r5A(v`{IPeRId|?o=Y7t-GoRUIft|NoVdw1@ z*m=7db}To+UQ;*1&g%`ZW4#`BURS}+>q@wv;=2xZp09)79h z9ot*5V|x>JY;VA>$=6}m3S72-LGHfkgg001iu(fyrwiat)*WB~4 zYwkJN`}VW2>*yKSd(YEwiFADm&a{owPr{*nooLdL+4DcQ`;xw{voB!R+2^pe{|vVF zpTgGu6WCgQ3|q^OU~BmyY%M>4t>yc$wR{iu`g|9TSZij1(q^&7&)r;CKhxK-{scSL zA7RJ(1MFD8haKy8u)anpeeUzpKF3v@&c>E5Y`&B5Xek zVb@OuIQ(jm=5C$ZXqsZPhD~7mR)E9ytuc0M+6cC$4PnRK0JiV-VedQjV8>Dyb}V(^ z{_?jrY|U%Iu9upy_mRJ-5$_{4uzPJ(hrN$fgRNCiwpQ{#rz@tfwQ3Dpt1V$`wFPXg zHixa%X0Y?VDQvBZU~9Dr>^v8etA62dIyT1MKLgs0U~AV3wstLHYt;hwp4%L@R?Xnh zs$!bMYiVn29(yY|(=ocYNo%Ecy4E$hGksl?JHgg@N7y=dfUR?T*t)iZt?Lf3YjS(o zI=6+rXKe>NwryeOp$+Ud1U-4%9DcY(v#M$x1r%Txa6SOe&5-}=M$tsm?>^o3pT`@zmjAK1F@3tQ*juyyVQ zdvE9oyZ(E?uK(_E|12olZm_lM3R}A_u(j(9Tf27Dte-JDVfXu$ec;ftvgXXRjDLsZ zSclNpu?~jq%b?uY1Jl^qJAxw3F$9F%O4R)TU!p_qa*z0F9?DaDV zc5EfE*Uv=Qu}*-U*YU9PIu3ST$HLC*7`VUU9}PRN2f)^H6zsh24?C|TVc!!*z+NN6 zVefmxVAsJ=IP%WYE;}6Cd~A-b6z-q;MKce!uXADhItR8dv*DPlLNpxX%A2yU3;rKl z9qKQxD@^@;-6F1=oz6$IFpX1o3GBXyV4JSK2g7s32f_2j2f_=)OW{T0CGZmQVmSP) zsB`nu_wh5RsXBiseP@cx|1sn`ZAxT2ERCD#X4~P(mT6-xj>NWD;~bIe>v5K0vwugW zzF92Rp$KsoYTV_yd5+G_<8fDDvyaE5zFBM@_t>1x>-@N!t&{vYKG`zg!<6MAevabx zb|QUU!zZL`Gd+E6{{MV9DfK0_{~p;q`JZJ(4SQd!4SQcZgR#7?oen$R(_rswr^3Dd zowcCg6xeZ`3`Z^2QlHGX{Ld0DqOX0q5VkKD!1m>QxQE6)5B5HBF6=$<9N2qctWoFS zY}ot8S+MuGGvSD3LpsO5NwM!&($_j&0o%9BVf%I&++VTOBQN&tQta0C64?9g#c=r6 zOk-p^=l{*hee`t=+zZ?HdtmE+H*DSSf*sqPu?#gnje@dWH#JPtb-kHP&F z=Nj0#cocRn9)X>UhhgX9A=tTi5VrQK;W?V`0oeD1`{DoEZ}LBDXqMai-r`zo|0Znh zo6z4{zk%IazYbg5*I;Y+D%@YbzXDsUmtkx55*)E)_axWk+sW2I<7G7#@9JCY{8#k# zed$Zs>;DVbHQ16k9rxkbU2~tq*5NbQHTEfN9X^3wV;{q=v5#Qa)Q50?`TqfIUEYUX zL+`<^nRnq*UH=XoI#x{Q?Ay24`m66Zu*dp3U0X+eGM)49*8iffb^H@{J^cY&*WY35 z`WtLre}(-S;1}4w{|sB#pJ3P9kFa(90q(CjzK6Y!eFuk58>*@NyYxEz;)H9cHtg7I z!H&Hq>{x5Sjiref%@`?)1-O}Bup>E^KWwHfT%*%S`H>ZZA^+0NKJ_fD|ud`H-I-T}5R?P2@U z4z@2lz}9ek*uJ%et?72K>v~(*8n=Oczug9QZEg))_pRWFy)eb@HL`22?=GotmOuNs zTe4N!fQEZwb3Qx5)^rcpn(m&)s-4F2y{$Jk=cgC!v3tVSzX$A^?G9W2Zm{+53OlYY zu=Vc@TmMe5_1_1!-h0E=bT2scZm@yaMqzVo`@@cHBy8PAz+S(@Vb|g?*jfyQU5i6t zYcm+OHiKaAJp*BDGXS{{#t`+l`A?Ao-hqjP-^NPU|nU5ip1 z%|dLhp#`vYo)25+QrJ4rgRS#i*!4OGw(hfG?@zN}=V2!7eRT%xyiA9kmuax`G8J~s zPJx}L$*}V@2_7IlN?_-0BJ8|PfStGTu=kX4uyZjMc3#K8k+*CwxAq5Pv-StU*8V`) z+AoEz?Go5pE`}Y;A~@#G_9)MNNU~LoE*zfsFl?UpP}rK6|03zSzv=05kD;&YV>xUe zkA|JkWx27Bf}Ou3Ve5Yc?02b$!x3Y(Xwoq^wecCqbDczA*UE{o=Q;s)tsD=#7LJ4M z%dxO?zXG=2L385^NBsU$*vgvtmLs#pXPEBgr^>#b&N$ZlG}wNg3fs?9VCVg0_`m#! z=OWuY&Kc>ttR|e-v$45G&VpSdXQpclWB7E;O>O+$LF;%<>YLT2$2m9IDrJMvX5+-q zmt0#H(ARs*`LOeSUK(pdKE=;h9Q%c-Z&pV|n*T~{UZ+>U*5Y#5eq9FJuS;RCflFY| zeKG7hx(IfToQJEDE!$5V?{(Ow%I{V1bn!}fhB*5BeZsZyBH6Ej7mKfk?bk%vu1~g1 zYx{8{HqU!Q>YLSBiN?7Fo3FhYo-O-L>DsJz!Zat@<9)yV8<);Ub8GIprLx_YY?<$E zWxE5LeZM_7j(xi`*|NA?GxuO~uI`3i19#=F?Ij-~kG>~H9)0b->DsdJ({-69`FG-v z(AR!F47-*eg1sIeguOSdhJ6ou0Jg99!`{2@gCi$(6i?<`{+;+*`r5bWVf*$R>~;1m z>~;1GY>l6Wz283tdmnugwvJE0-bWvYy>C4R``*6>E>&EQ!mdB(?1f}&p#GT-`FGYI z($_kC09%LmVe9Z7>>RxdTbFlW=jd(NzP$xow>M$y_6BU-UWcvQYp`{C7507d71+AH z412$O3HCkxML2ZK_F~WVCAI<5@C(>E`y93}pTYL!Q`o+I0^665VL!ip1cxuaA7y!q z=L_fQ8~Qp=U&CHYU!}1s%l~Ywcs?lB`5$uEeh&|j{X5ukehY`6O;R6^RgI$aJ)|mZ zzpKFZyE1IQE5WX}im>ag5cZx>0e1Zrz^=ca@YT=!|G=)tzhT$oU$FO-KVjGBAF%85 zcX*(*`VF>LzrwDGUtrh3&+t@TAO8Ah&OgGTbzSwz@|b^5x^30qcl}Z5zVY zwgGHy>%-Qz9&Byv!q%=1Z0%~p)~*)p^;Hx0eX|BUTX9y0BlbEvm&Km{+d!Mr*RdAC zj%5?rzHSWL*NtH3wH0h%Tfzh7M+?|_ZVo%o&0y!bDeSf01TNKhjp2y3Vmjwow!mgz zHi!NBeY14!hMrU6cW7hWvL=n4X_$X+xi#Z>t!@R|k7C$swKeP*wuED@Y|nAuw%A-t z+riGqwyBM(e42|7rVn=Q@g>|dRI8KZmuyh zt@FRjIFY{Ab^>f|$HUfk9BggJ!q#>SY;8xw*7g9{+Kz&)?f$T}9SK|85wNu#4!fp@ z!Pa^xY^{gD)_O2JP&y5Qt@S|IS`UD|4*J8c`+l(ZpuTX4#@`Q)+?3tE()lbm`QP1~ zPhaPv6m~A=!Pa^%Y^~?O)^;{*EoZ^jawcpoXTa8OI&AHx!2=cBRM^^0fvw$S*xF2j z{mfYchc<;8Gt(yC6%5e%1*vbdbgoFYx5(@ePxGrJsq2MI}NsOr^43l6xeZ|3_H$~VEc6<9DY?Q zpStYp?DTzn?wz3XyW`+Q@tItk^^MbKq_MIX7pU)9Iomkd&Q7+9>0GgF=V8n0J5HYq z&zJq2bZxdCeC-9;ti}1MZ|0wUzc6R>92ezmGc^Cj$%a&Ahu6%N*t}-0fSt?BVdwHP z*txtEc1>RbN3Lp=PgV9ciy{Ax;(Gde?p3hY7K8#zh>~^log9>n_;)#GSDBg*#wtcsuMkZ-ed2t?r(Z<368i~DrHBf)JC%!n`3zZ_8xUVY+vq!2kQEJVQYR5>^-Mr z*`({9p-!DZ z&nd6N;alb8n>GCeo7d3Cu;=;+_FNyrUQ-{yF*e&q*-@O@Xuif~9lnCC!OC5bX8$%RRy+Im0|A@m0;IQMc6e{2wTevu?G+rnRU=td6ZF z>{x2R_PshBV`a7B8m*nPjZ5=IQzvJeDqFu)3)z@iT}GF>x$8V;y_~H?W7bc$TIpP| zY#U+gByI(}o?60=s|D=(Y7V=;n!&E0rm*Xy3G93~h6ic>MzHhS5Vk%I;LtqVqa4d- z*zEhJu;(g*T`QZw{#>>(9AlTIfbXo!*4P~9RwukLoTewFy5$$%c*Z8)ubJPZQj<$irm)dCz`?4oC?>~JJ%f5BQ zZr}ER?c46K=iUwWn%fn&FT214lONIS3=b0T1Y5@)VduC5?ERoU9I=%(eCOEubJDT( zgB@F6*s<*gdoA>V9qYcZ^Uxc1?7d*e+7ot7^?+Sd-C^%P-C);PSJ+y0fd@&e&ak!W z1Y4_pU~9EE?7ed@IJB#h#&j%0vH3nR1a@8r!;W{tiFj&%SWb5%@pIZq?7S@+?v z>w8$bwrRRHnpx@l{NLM|O<#{a3w906gk1wOVAsHO*g2mDJLgkj`!)r(Zj)i_HVL+F zC9rjy2wS%auxnyG?3x${yC%lMgQUwC*flX4b`2Z=y9P$V-ZS@yOEum|*m<*d^RQW~ zx$q!eJ11RRJ)Lu24#eiXEQLMx64-Gsh8_1J*l{j|UHc1Q$2}i*+@)~DwW0m_2yFK4 zaM->b2D@esh3#u>M)vg(*mEBY`#Jp}IDBF2EjwHvE3kR&V_=WH9JUTe!`5LLY#ok* zJ@=7t%$@B4&e2KO9QTQ^YyAY+wR=45HGdpz-;af3t_sm`SmTqEt!BAty)|49oBcQk zw*F_sUXy3R_HzKjdJUe5-F}||+t1TsYkV3!DESr5sj%~T3TzGS>qXeSMlOWCzAu11 z*ZJ_E7%m;>!OriwaLiRZ*`2SI*sSSwuzk4}wrH415nEw29QOVC+_kGx-|V^3I^2TIaor3%uA5-{bt7!QZh*s& z#?f$i?8mWL>S(=ySDFxy(aI4 zouhkTYjHPhE$)H`mCI1}^-kFJcn9qL=62Zo&26yP%B^r{TQ`kqZC}M^Enk7H<;$?O zdkMC7FT&R91=w1xg{{@|u(f#(wieIA{@Kklu(f#_9+cvU<|)|PJPAAYCt%mNV}Bi+ zV|y+2&3+bI6qlys2W-Cfd)PJk9c&%Hg{|W^uxs*b*mHjcTh}jP*W?$lb^aXoGtXzR zb^jE$?w`P}(T`#0;Um~J`XTJRd;mKS@56&KacJLzoriZ}*XTR2YxHf{d3p=>Gtrx{ z|2@SUu=8jwe#-UzG4-vG{PR3NCtJm2bB=z+HYlBs<`;N^cnJ5j@FBalt!#hLcbxcl zc((XA*#7(4zp$04@1LpfhUO_~wEjH*(6?0A{++JP_H);4A=fTY-wJT}l<71>wu;G? zjpNuVV{>ek{-3^8v5nWbRZ`zH;!#Ju+7sr)pO%aldWd5RZi!+$W{-V*HvBE z>!J?qx~~ln*7dbu=h0)=$2LfPqkf`qS=sQNbJ2{xzP2gs{h$eKKO4iY`9`qUe?!=M zHGpI8TG4QL?k%u+?#*GZgUw*iy(#Ryr3kiWo50>XHioU?MzD2h1rIKFtL$q_*n2|@ zxK#G$ur;?Yt+9FREmPmJhN~R+?ws_Vy&G&Dc7?6OF0l9Honh;;6YQMq2s$U@I-L{8)FKi2Y58e)T?AyYFrCS@=d+;`}b>15Gp1T$7J+~MR-J7H_tz~~~ z*0LXLE&IaOazEJh-3PX|`@+_;H*76?!Pc@T?E6~}*xGi7eShl)TkEc{weA92+s?4H z?F0{&w)?==c5m3)?gjfZ-=46w?g(4!J>bx~cABHL<{E{~+Q*)6?f1uS?MK4aegtgo zhr`xx7;No^!q#R8Y;6X^)@BeqSpE)#t<3;9Vylq)c&-DItwpjq?uFQj)A?u?z}8_t zY#mBr>o5#?xsS^*E%T*tu9*>c$X%F%Gl-5_0K-7mmq-Oq=u`+2Z+KNt4; zI0tqf&W4?bQ3UMuaTa#xSJxL!)7h6g{{Rku(h}vwiZ{xj`d2|v0VW>w##A1b{Xv0E`=T2CGcRyb}{T4 zya+Cp{X#fmt&+~UE^fl+*l&bg7dOC;^?KN`u7aKSm2k{eE6vq9wb9&*&2v|VJ@-AZ z=e`?u9`1r&Q+LAF`3~4MbvtZbZ-cGtt*~{y1s*JaZ-!kX_I)(&Wx=6M))tPjDisRv=#)N0uGmj_^LaX(z5Snq=))_UojV|^W)V|@*F ztgph((<`uR=4IHizXUth7h%Wx0_<4V!jAQM*s(qbJGN(G=lL1fd43umtk|D|t;LgY zXpueZ2F;8w9Ql89@D+Wn)t9if8b_U2t1qy-E~>%S>T}pyeFnQuK83B_C$P2q7`ApF z!Cn&|!q)Nw*jm01Tg&%gKZm>vTibWw!R6AFeSI7DbH-b+?@@2UQ`P?sIJ9n-&RNUK zjn~(*5^OCi!mhhQ*xFWrt!)8pZU13zYxy^9E&qb8-Jh_v`vbOizr$X0zroh>SJ+zq z0$a18l9ohkcLx4)*7aZ(-jjzJYz8_!^E}lr?f08Z>wG)dT5k(m>o%~p-UhbTTf^3RE7)2W!`8Yr zY^}G1t@ReLwcZ@I)|S`s_OP{R2Rrs1V8^~a9KIJO z-`l1(nmw>h6z`t;R?xW_>HGZOSD3Co9dp-NpFMN7ak_4=oUKH*y>m8Sx6l7$>y)!C z(U_f+tx`JYHPIcL*F-niYoaUcymm=rWxB?YnO5=d2Asnlsc-!>uk{&@&1-iU?6o@- zwx2^_`#Bi4pMzlgIS_XJ4S>Bi`@^okez1M-3%maIgI#}pV6Vk}Vb@=8cu0ynnqIKi zUQamcvq>}@)@CL)Ycm73Hq&8iGYz&jQ(B<3|p&7u(cWid##qh)@~wfttPqQ_gL84je)(EN5j_g0C-5bOl4n3!Patr*!$Q>c&hqGJfU?hoy%%A|2Jb7 z)7M%rg01yJ*jg@t9s7LPv6sS*bsp?k=faM44(!-w!`_o;!4XS^4a5?EPa$j#H(<;E zzS#=Kam^hAyXKa|*5PQ_IxK@-b4S6hxg%lgbOh|0I~;b+9R_PJpe`@vwC| z4i3NSrMaWYYB&EkTQ8=s=e`Jb&Mt)Q%LTA~IUlw!=fU>nT-Z4|2X?KW4SSuR1v>|4 z!p^}N@DRmyI_!0P8tgor3jdcD@i$q0k6%S!=V>MEJY5G{t7~Cvbq#E-W|4eraW!_= z?@a8j->a}&t1Dq^bp`Bty&SeSm%-i#E`_blC2(j_Bbsz%b&>!7%5I~t^Li`n{rnc# zzTXVn_nTn*ej{w(Z-AZW*=XTduE!p}7A8NPhdZ$OUUPfuTWbS-AI9dKJp?;v55lhJ z)v$B)0PMK#hrNf~2RqJtVaIh3>>S(;yFTxNy;t7}|5v?^rCwX77^8WCzOL7`uyuMK zcI?l=j{RBKu|5Mk)~8|bpHIP#^-0({eFApukHgOCWAG69zXo=_J_?u0{s?R>9P0m71-Lm47)B~ zg001iaA;8>8V=|A!(^*dZd$J}n)sGuMtM^Ze>cPHqe_$Ye~wuhb`4d6txZMP z+7!aprUGnj3SevV4^DYK{S8~2zhG*~{VAs^Iu(kUIc1`^Z4@ufY z^AkK&{3GmlR6oF>ZC&-rw9WtBi$=t5ts25!e+^)3RUh`AT@SW4bzy5$2X@`mhOJF4 z*jm(ttwjxZsN$#&dmpX_PnEqY9I=-*d{-2hj%_&Uy4f0b-E0M0i(=Sy(;BumTf)|6 z3)p+p=CJE(GuU;tDQqo^U~9PvY%Mp2t>s3r>#h}SZCk?Dwgqf$o5QZNX7JEVEZU~9 zwQd6Y9@iKStt+Q%?ECiE>}y-tzHJA4?rmX@-3E@a3)9%8sf}g_Z1csXJaf3O>!Ce; zm#A;MG;YOoZk%j8VVfx4F*i?#bZvzVT)T75R-$osNw!Mq+;rJ?$2Lp68$4UQgxD5{ zcg@W+N47n3Hv7;q*)lzC+bcKco~duQ3-X`LGW>PrP>;E7NRubvObT#d`8D*Kw%d;G1} z);fPYeT&7%!Pfa$*gCI(U4zHKuEFK7Yw&2;d07TK4@bey!;!G_a0Ki;91gn%4})D} zhr+J0c<`{6hrrhIV0ftFKM1zA2g3gRuoV6;ZS()%*kU$2Yj+XXTDuEjYj*)`?aqf? zC+ESglXGEfdk*ZHKO6RXKMS_@XTsM04A|PA4qN-vV6TBwVQYH|Y;8}5t>sCu>-0p} zYvcqtw5=aaI1!=l!PaUeY;CTC9s9MgW4#7;tXIRXv#VgoekJU>y8?FI zT@G8D%V290&kfe*QrLBN32bdHhC`cbNf)ot+p)QBZi5~Bt*~Ri1$N$VhONa-@X$0* zG&jQb?FRV2@)Ym3?b{mqI!}+nuK7n``~EQO8h!}2uMfiZbv5j~Jpemz_ruQHeX#R% zFYG+s13OQ5!$Tw7ble3yZ+F7h<_)Ki}LT_f+q-uK>vT_f+p-sj$dts8pXzXE&h zz6^&J4RlSWMZ7EbdijyQ*5(J;+I$aNo9|$+pN_EW=3Cg>eFM9WzJ{&USFpAE61G-f zz^At=&g(XjMIFRV>Z^!e(v$gssgVu(kOewl=@Pj^$U_ zzWxH+*Pmhg`V$o+&bqL3RtI*@YQxT1E!g{cO?a5%ssX#+ycRdX=6c&0c8)iKty?SDzO{tyTMO8} zHHYn6GkBPMYYK;Njng^%wj(zC)&aI}?P0HlcCgpN4zTx$?P2G&E$qB*2RpCZ!p>tG z*m>Lrb{@BeoyV`jFUu(TLY_0c#t@WO;weAR8>pfua$-BeWemB@_W>?sCwhL_Sc80CpPH@E5AjM|i z2V%4D17Q2!AGYuPVEft^wy*oa_O%b}y4V+XUG#>%kM)8bTTj?^(gSv^-C@Vt4R)P$ zgl=TWe=+aLCt8VOsw5wNuz4qLln zu(cZsdyNf&t<_-IS`C6D*0Qai?}`G`u@IYUegSMP=EK&a6t))gU~4fKwia_>YcU&k z9nFHR%}m()%?#LDO^2=3G}u~Ag{{>T*ji17U3ZgUYgYnWyNR&tYy#{$8xOnA#=)U= z<)p8*JP@0;TMFCvC9r*84BOX5aEx6cjXg89(HxYsjg#%*oGnb#l0E;sm9ur-A-U^3 zX8C`-zRs5tU3X~ix=FGfma~N^^CjLjM|>5dNk?XzgUxku1mjqn!_&A`({H4TmUrZ_Kw{Ty^L{9iSi|NYB5=xZ%*hn>gU zV8?nZ+*@^<~G*uLHXhp*W_7nJ2K`~AvFaWWm5E#5&nCwDQ9{XV+t z`hMS;o5wNSoov;!Noccq;vJ3G#e?+qysLBbJ^;J!?}vR)yAQTr_rlKga{RV__h1iw zYA4_9%d^^1iS>{!;qj^%mS`^|H3_)2>IJ+c~!-~Y4k zalf>$Z^2$eZ^GUW-he&V>#%iNQ8wwizg-*-WOK(m3G4JO*E*-~+!kY zV|_|r$MFg5eeq-1_3{yH4L^ik4s)?|cv2m+#=1YeV`jCv07hr?%|-aj@6TAN04*zr)t`H`qG<3R}lt zVDAAx!x3w?p9N*Q#_ync9o5RkS`&7xHDJe99qthWrlT6{SgXRlGb3#k*fmobc5Ibk z=cyv>SPJ1`8m|HzvDMJIEVlR^uwtEWOkc;^2=>}*2s_pWuw$zadp*>HowvHM^HK-) zyTaOV_*$5J^4LwWS+^#sZ_CuTD7DeF$L2WN!H#1G*l}zR`<~PmcD-%~JEz;i_N@)< zwYLpyU$=&x!VcQI_;TEo_DOW1373)uC)IqaI*4EEaI6m|_2!LFfAVCQ9H*z08@ z*frJ)c8#@!hfDt!u-|nwhn?qUaO64LQ=OaLu{jsJ!Pah9*xKy^JNBJn$GQ{j*mi^+ zO9$9BH!d(89kET#fOe1nr|(|aeBanJ^{tSuw{7p-I1@BSCu~c_`=q{A)43Aax?{6{ zXV7W7?5D%C#og#%D((u~uP(6X?hJ=E)zdZOb?tuG?DuNsnkai8?32a&!c)X&(|?Az zH}-krUa)=b3ES5muxltN(;_}kcj()Ipy8vJ9TVL$<<$Ubn zOT}ye+HAh~OyIEwbL|}U8w6Y9pzPWux&3*6`nHV@>9~k_y-$q9?phxKdz}u4ox@?U z>t!fx&4<9&*z0V3Zme;z^M5J6P0Adl9gE$XkAbb}XxMWf0MFO;qu}tRQab0cO0YS9 z6Jcw7CBAs<3AwSX$z*Jq{xQ``uJzcHusiqHVR!Cb?>(y+*L5=^ zckT3aZMMHmj-k>qGubk&;&c{MJWtQcI?`UBn`buHTK8FLtjfA38!P_(N++G4Ltoe9 zT-Y&%-E-fRv@2^Er@O{llpAYdZmb2c>uWxo`4z+5!l0h36ub4D2U~mB?$X>?OJLW) z;&knXG>PX$-+vCFulLD=b7LO_+ph!D*qJ8Q?}XitLud~hP{{ow@I_u%81&;@gRS$m zusyCrPb5etkqMnwR#eEUY>xh z)#I?WdJMK!YhY{hDC`>S4O^>6U~Ba-?7iS2*xEe^Tf5b;wR-@zcK5^9?mpOB-3y0S zjiO0MR>S#s;_uMc+Pw{1ySHF#_a^MRdjqzXufx{zHP~9d3cK!Jfvx4su(f;%ww5o# z)-s-p{Y{TX#FY;B*1t?hHLwS5-$e(($&S~kgqq|LO<|8D3H^tG1X!`AXU*jj!I zTgz`?Yxy;7Ex&@T<(IIv`vSIhpTpMfGuYY=j_gUI07xe~8U{_uuev^i6eyAEo@6Dgy3yCZb%pI)7udetSatmzcgF5|>;!u)?gP6X_l8}Md%@1jo^beHJ(_f6 zzUSX5^r5fwvM=mcPMP_{s&vTgK)w>k2wUpwLBQMmIuMsat(WnYxY3w)^;iE&jm|h zYrPn@){9`*@Iu)4z6G#zF&}m=N@4$OY98!d%!L>Jt2*zz;K)V8bk22r1~%v7blAB# z4R$V0g}6{R-Gx9|K$K<*>Cr8n)I^ zpw@aBc58hU9NJ~im)7obY}W2F*xFqRTf0kOYj-hh?Jk0?-G#8Vy8yNp=fl?GJlI;q z`%G(bE_Q2i4s0#XhONa}u(dc79-e9_nq_5+d)@Dx*|#g`8@e@(hNGzLYx$RZv02A^ zVC#4{Y#r}{t>c}rb-V+%j<>_s@iy4|!L6|GVYk4p!JA>%;7zdigBxMTdIRh{Tn~Fc zSOq&TD`Ds5I@o!+7It2)ft{DDVdv#4c!c!55)SPP(>ZH-KQ_m9U+P;S^({$lG*44M zF;1muII{gazJG}8X2&<_co3eK0qtsdk@$f$Pv%pYrX+j(KDFn4IQ7kZnx%0b$=Sxq z_Gq$IOy_KS9NPl*eGFbAUX!lP=Jd5ssy5n<(Z?i~R3c|4Co(jeo$_=6Be6_zm{n_$%zZ`~q9cpJC_eC)oGr zA7N|x18gn7hppvzu=mPuVQc#hY;C`Wt?gH^wfz#dwrkmwtmPNjM@YxdVLuyv28Y&l z)hE+B-d9_T>MhpSq8jYjtHO@G3hdY`!(RV?Gq+=}gx#@LgdJ-k>{u(nuA>6jvH!#P z-h*SlQpNdKiapy8!<5AyzZ32?SCg?_M>XKkt#&l&$gYid)5Y0I+NSjNSWRHpXye@2 zjbPVkL)bOg0QP#W4_o(ouyw5qyQb>E*10xpU2DM+8>Lisc&?4HIp-U}p1T$7oVSEs z?=4{0Ky%o>G=pRAdgW7>ea&>u|6Sbe=lUzM-5hqTo59Y@rm*u;1Y3(u;LxIa(!wp^#clz45Zm@mn3VYA$lEyAf<67gM z*hXX%()LJwvvp}r`(d*mePL_3AMD!h1G_%5;H?905DF;c)oAq5DYwck*Y_*Rjli9m{msu}p&<%T(AkFa>sO6JYNVlVPu` zNw8xrfgS5acxr@`ju@Q>leb-Y5{Dm=EJUwQrLAd5B4537xo&M1J97I zvtd8$%z{JvhU$}PpZ_-kr_tAI=2Y0ZI0betPKK@hNwBp)5w`Xxz}EVB*n8x0u(duG zw$>|PYkLfAZI{E=_GsAJE`zP*QLwc<684^a1RUB`NZL8J)03_2)eiQe)F+yi*sSAq zuywo^wvN}puIZ~`>v|RJxvzw+>lLtT`f}JhUk1CTFNLl9C9rkB7`E;g!LIQOVdvoj z*m*b~_IsW4VCUsr*m*ezcFms+J1=L!&dZsw_x&^A$P-Jm?6CH?V6*l&!`AX9*xKC) zTbmnTYjHj7SXRL?cct>F%f3!d-^cGh&eHkY=sQPzE4)DbepvwP-ba-6=0_FB5OR@Q+wyEc9gxLD^uqOW885biAd2eAEoANKz99_;<+UD$i= zJFqo-8}?rM7VLZ6o3Qi#1{^xnPCnV!@3Gm}?_g{EEo@)Efqh^38n*9W!PfCh*uH-O z+xO35-v>T}9m}V%*V8Ak?+YKpu1oJnKO|d|6n8XP4)gC-tMU&ItwR;qI#h;zzpey( zy;Ovq<3iZJ{Dd#osRDN2e+yvWbN->fYv6C#y8Z=Q*FRzF`Uh-Xe}}E>Z?JX!6}GOw zz{3GA9}3_BN% zVCSMC>|8W}os0Ug_mz6Eb5j>~ZtB3!O>Nk@sRcVXHDTwb2JGBahn%22;op*w*^Nz4c$*4>o^R*gN&Dmiqes)B~G!{DOR1$L`pzLpRuSb%h;g7ua!jh8<@oIQ*`c ze2FHjmv{&5dKpSz`#J=+uY+OxItX^o2f|)!17O!cf7mt94|b0G!p`x2uw(B7JNA8H z$KD%u?7iSp#oZGQEo!B6jy0aU9qU;7I@U3;>-Agm=6zx`cIWK?*s+a*?fd?)eIE(i z_Ytu7%;B)>a~SM;w6Ei_S@&_NZ_N~U>(oXw6Px3n0eh_JuyZ^Ob`Gb))?o_lxF^G2 z85tKJ5IL!hU|42m3xS7Y<#rJtk;!eBrpf%KEV!N?*qk zpV935|FAo!hroW`I2d+L;xn6TGyBYz=^Q>}zU1Hg9>-X&w_{=Jz5=$d$H4Y=IqY?K zH0*oaGT6Ew1zYzcVe5VbY~2rs9ou1W=-e=IaZ0EwR zsdHe*dN%AloCQ1fGhu&rIRp0kI30FPo(5a1Q(^D7r@+?kWZ2rB1Y5fkVb{e8u(dlL z{x9ur#8xaVZ-A}k^{};E1zXFNu(iAnwwCev-ZgzKc58VJY%Q;by^mc5Tgxk9Yk381 ztuBYH)n%}?x)ioHm%!f7FNVFoFM_@AS({t1dF-2E*YHj0+NMe0qSQw7BsP!r1ZzB zhMlLoU~73NY%TAAUE{aI&f{&c_vu^V$XlgoIIQKf*c{6d=AZ((zdz6m?Wufvly)*J9_ao8g#Sv-#Y?VQcE`A)J`(^#1fJ5^o( zZ%2Gg-|_1A5j;`+A)H+sr$2zF%Kkn)L;N0WKi`EThD;yV)fd>jMn8wW20w$n#y*9; z#y){#Y(gqKJodNPtjRa9*IU$`Yxrx}>+44b^?LdWyVui~aLm0SjpLnWak)!nUw@&m zeft^qUi1^(UHyNA?fVa~*Vp&3_u%i~@TEo?!~0}~R_o7Q09*5axYoY>4SV1F3$`zR z!rsUJfL+tS!}j$z*uMS>hp$Lfb~xWv@y&Brft~-#u=k-#uya@u_MTP<$6O7|r&=#3 zEwEXq=CE_r47N^9Ve8ZccFr2Z)~OL}-y6cN-3G8@sSjJ%da&!YE^M9az>d8(?E0+* zI}bHs*K-Zn^;{kH{h=Bh+BHhrMKdXVpMQ_IEq%SuwSleWHn6qZ8n$*@!LEa1*mckv zwzgZs)^ZEjT5b+o%gtbG`3q~?T5gKnS{A|9aue8EZVX$?jbLlp3bvLl;n1>9(l44! z%lv!7UFmBrcY&?t&amrpC)ir<2wTezu(fCpJJxovW8DFEtlPtmE$)?$tu1!1x$WSH zHM?hdjqR3f^`lEVG9B{o?RwGII`o9CLl4+GbcY>xH`q1P74~}U0=tGf!>*xDuyxx9 zcFpV!TgSa%>$oRu9XrCVnLS|Zx;t#$oYUUPwxRufFgEKn2)0fGVXwUbuyyGVJMMn4 z*!!?*5)ACYyCjju`Pw|%M#dmSq#Tqh2>MNr{n*ydF(?{-`c6K{W=z# z{aOJ#XUD+K*>c$Re>ChpZW-(x9tFGpkA&^p5wLT3IPA3^pEtag55w*p914dH)iw~@ zso3o6DX{D5WZ0UY1lyMrVb{wEuyb-e?ED`GhcDTlXdTbTW?#;O?aR5aeK`lVFK5H9 z)w5vxawcq@&Va4c>2PV%D4NsY@O4AaQt|hti^H6b8|dqrxgK_Gt6;5v>+Fc4;yGvkecQI`3E`pus3t?+{0UTOZ zN_si=Td+Cyn_>HQ6Kr2@Ok>qeW3^6gG!J2O9v+09ht;t2@Br-n^?umC-3QyZdtv9{ z9@uO5ZrFLa3w9pvguVaX0Xx>)VdvpC*t*{ehmH-S;cyMUjLosX1bg3l5w;dDz}8|d z?AV`&9s6^zV}BNQ4L$=q_NQUT`V{OMd=hr-Pr#1-aoBl(4EDWd4Qy>5g+W5U*Yxeoo3OQc1GYA=!`9|C z*xI}bTbozlh;>8HAn`8Av3^fq*U5LV>*QP5T6_b0&3p}8o3CJN^CfI;zJRSwynFIJ zIhqgQ(587b97Sbc%fHmzXnk#Jz}BWZY;CH+ z)}|_KZK}Z5rZVg`RSEV!UlF!;g|PRR3b3^-fUV^}#O6Bt8+M)j1-q{PgsuG_u(kgk zw)Vfl*8W%6TK@uj&HW7f+4(2f&uBlwBc=ZjaO9?9I%lox`d#+k=Tw6I^8%;_2 zK7NmKoX$6(?-X%;c!9VcyhL0#ja?y~n=V_!oNbP5jgl>k$38U0W*?fs_MvgQwlIz7 zYnx-6rFoj=`p%QBB{tW8i`18GpzJ7C-_5Xjo=xGYvKPU##hbt$dt-RM>>I&L#jW7T z5$Py9JXdRM6Vv%F$XYU#NU~`S^o%&WtzS-6(**4^VS8Udz3+(-&bMD$0K8rU#E7;%e^mQ$EOV?IT zK3TWk*e0d((e#3?TTj?y_kcb2)@6mUZcm(|vHH^2I`0QNzuV>V+Xwq1UB55vd=D_x~KK1NpquJQ|YW!KSeV++i=NYhLnGXBer*3Y&Ov4`ak-h5( z`nNvl>%1IDU&p!>?i?4SV+ri~Sqxi?MX-Hu!1(rcA$I$^0QPg)d^mh*$6&K> z%VGO=G;H6N!S?Ma*fn(|+%wG+%@MHkb~x;3-p0h@JRXMKc{>z#4aPHx^Y%a3S{wqq z1{*TIYw%#~*6JX*RQ|>@jVm$Q5vCwW0NV1vYDOIqcXkgIy<= z!mf);V1E|67{zacT{Ej-e+FC$N37ZX%Qf;4Hm{d4@(HL!hs6t=IA!1ncF z*fnI|Ud85p>J`|&?!rD}-CxFT-(P~AhZkYT@&at#*TP;8&%=)GIoPp13p=)FV8`|} z?3#KCj#w*2!{K^=Gk5J9uygo&ZmieRSlNBtI(>xAI(-ORrw?H3@;+=`-h-XvcVX-F z4s4y?hONWq~FTXop6RfAn4 zRbkgi71;Gz8MYRcU~5qk4lQaYEgbts*qp~!u=CgwcHVlSv)6SC?0$A>4qK~cua4g%>*RiyLy)SG7JGQN1-($Ce9a}N% zI%o~M4&pt7W8D(F>tGAmv2PAL_RV0|!KSeDUIhC&ZWGvburVCkG)dY-ldXYxk5R1i zd(+o-u@`JD_k^uwN7!p*57_%&2V%F@yJPqLp*?nMzZ-UIzbovz+6A`uJHys~C)nEW z2wVFOu(fXwTl;pfwci1@_S?hOzAbF+w}V6L>|SbLdttM0Jz=kj9u3~g?e>RVM3)1;$N^^bV+{F08 z5$`Xv?}Hblz7+#F#+5f^Ul(J`_RBc22%akY!rVN*wg>y0`!31#jbSoB;@^Ud*SJel z-^%G+TiO1X^WhM9mh1<^v&9F&rQ!qO&=L2_4%hoIPG&xZF?f1;x3aH8;hEwg@Eq|G zj6Gj`IPAO}2K#4bhbCV#?Y&l)VY8k`!Oq2z>Do*q*X^;nYggp1JqC7emc!QUXxO#p zIvh>ROBCaAsc#m8eLfMJeLf*K&+)l=x@p{#v3abMQr}9cPcQW?!6)}UmA;Pe6xg~? zpugjD?5ES$IXf-Qo2@&~b0)Th^5cxuw`w}){on#@j_-Wf>;F91d)K+J_pWnb?>%S3 zzNei9M@(5g1ZA}xYu@|zCG@p^7sK}DBG_|Zn8q%XEe7Vu#?Jpgr2EjxabA_1>q^-7 zfGc3feL3v7E`v+Tr!4z=DI9UuOy@eKHk#P4Jok;T9)YdR!?2%2 z9)hjagK(+zSPh3(&2%oSqxiQ#*6vgK7OUSUu(kXccAb3$TiXv|Yx@CgE#HT&-FvW~ z`QC-Ci%}dkeO+mJm`|OLb_oEl! zsnT;l)@kITPC93;f1wk0D zv1|rAmQ7*r+eNVNX`8@~bz|7EZ3KHgw1Qn1En&yr0`@zR=CJFa8EkEu!q%n^0B`4sDu8la6c+#J|1rKHZVN)@~2j+U*WoyWL=Gw<~Pzc7d(k&ak!H33i?A2z#&T z09(uUu=kdBu-DiQu-Dl3u(fUrTifkmYr8FMZQH=sb{p8*ZVkH*w}QP-7sDl)=(Me2 zuemMZ$V~(F$#N6_mdm;6Ltp1+6nnmNu`hP#qBrbZ^n#s>p0IP#1Ge_vVb^sx*tzHm zI~QGG*Li2y+INDj{XVd@-y62}d%@OnPdKzfma@aK55#8Q2f+5VKkPc`2iv#4u=kh! zV9ynltiD9oJ@e7lku?k#c{B8 z9Sd8VF|g}rG;G}ufUVsq*xKz6hgLPB;jlLIv00l^*xJm4t<7B6+RTBi&1~59ISaNn zGhu5n19t4wVdvGp#ybl8dL(`A>k+W`iNj(0ei&@u4~4x?{15iNa0u)?9}GK|gJAE; z2f~hRDeTymz|QMp*m+$9JC6%t=W#si&FgCccCTe?wH%wZIU2S$%V5VoiLsr>qp*AJ z$2*{yt6?-8p6eX?IxlC#p8G7=x}OPq?lWM|eLC#9kAYoNr@_w4sj%~M3hecFGVDB^ z1UpYB!p_qPu=8|0>^vO@J5R^L&eIBbYPn2hU&H^TkFeUYketfZ7+eX?ZvRQy$H6p7sA%|0@%+b z=fkt5+j+3pz`1Z}TU&iHZS()<#GUlDwv(x8Yk3EDYk50tZEl0D#jUWlxCM5sH^Ywg zCfKpv2>YIP1MKI5>*0v4FvVuy@5bhPz+I_t_AFeKu8HPFY`*pd*fq5lwhqt3-YcGi zt;@5pYwQ`=HTE>@8hZ-1j!(kgXPu}D(pC3fy3`g<*ux!>HFAbrSs9e2hS6~3wxculg6r$`jn{e2RYjk z**?tKJjX}LmiZT#O-R?q&%9i7pVD`n#{C3dApSUwRWY5je_vpmuD+kcj{UQAZMII^ zs_$3WW{AH`eKQ^G!`HcSChEFxvCS8Mllo?{+ut9txu$-AJ>B8~N38mp$xWjY>4 zA$l$RMPILlKVj?o2b^6S!~G7=)VN17o_+fb`yAPSg&pTFuzmj-UabBJ}%`(XeR!J*)xjJ*z$(V^xfXqhEPb_O)@*`uBy{ z!v@LT2zy+cY12ivCdpPmohz2DHMUOTEn(kZw}3s?=5SZ_-wbw)o5I~?FM_?VZvuO7 z+Zgt~K8t#B+#6x{{je20uwdQM61EO4VAoJ{*lWHSJfdLT(G(8d8m&9^Z~pQ>ciWA= z)@E1O+Ux>bo1I~6vlDDBc7&}(2iW)5_OP{S2lp*lckBRLtLv<_bFI+bc3CTuCVja1$G`f!_Gq|*m>9ob{_VIU88%!&cmLt z_w0_a^RNdTx;Kj^9TQTU|NDOvi`Lh20&FeE!`5;fY%Rya)^ZGNEl0yW3)US6z}9vY zyze@5!Tzwd9tm6P5wNu$4qNMCu(civTk9dPwH^$+<_E#nejse^2f(h|{;+e=51v}E z?&u5u*M5@!`+*DS>$;u~I~S#}wU1{xYd;UWwVw-H`#G@pli9Gfo&{U$MXXV4JrldN zodH|h>9Dn(23y;yu(h26Tg%C?wVVW(7OXo;;LsAc{ynm~&i~u1gNoMI@<7;HE`?p! zOJK*o7#Vdwcc*m*t{ z_8MCOJMYK9&iiuMu^$aP@5^9oaTM%2I1(OEu*QcKw672iN2l_p?CYsT(QJ6{n_Ol6zSq##>+NdTdASO9Uao|lmn&fB<#O10xeRt5 zE`^X7A?&^z(UI}c~W&cj)-^Kd5YJe&cK zELeA(4tu?x2LD$s^8Xk4YWg}C55UgF{jhU!AM9M*3p*G0z|O_pu-C|4uyb=K?A+V| zI~TXZ&c$u8wZ9d%_P4-Z6F0-o#Z9ntaU<+p+yFZl*Tc@mD%iPL348Cn4jxgk?zk3? zTr^&H>fdLv{Le67q_5Z13$Syu7Ito)hn<_}VCN=2!#g+6Vs~zyft{PDVdv&4*z4p; z*tvKDw)T(1*7`BnTCahv^`o%0egwAG55v~>A-J?)-SHqC+SXim>R)Z+|8lpMZ`0S> zy#-slH(_h{2JAX~9d<0Q!S?l4*mdy=Y~NpoT^BFG;d_Ho4cqsx zVCVTu*m?c}cAh_no!8G`$NDMkJbwav|N9vB8vF=$9efD84nBab-TScX;62!N@GfjE z-+}*Y4aK`9Ygcj8^|dO5txW~kS`@&p`G1(p`_bR9wfYOTR)50QD&9?5n?JByo8Mt; z@f&O{eub^YFR<(EXV}{O1iQ|Dgk5Jpz`j>}50@6KJHCTmhu&`*=d^4DyY3po*0urc zx~mUc>w2)Yt_xe+I#DG|tpZ!i%CNPp1cz4j zqT%RN-jsda0-Lql9QK;s47PTg!rpI+U~9PvY%Mp2t=&ejwQB`iyJgg&wP}gn+O&YJ zO>@}VG=p7-O<~tv6F9VJA`K>_@ALn5bvOE2n_XdRvkPo(c80CZPO!Dv5q4d5fUQk? z*lV~QY;AUcU02(~uB*1N*YI|*wcHl=Ge#TO+HM0|+pS@1yA^D0i(%JUYuIacOW1W6 zlxdy+|G&D^*IIUitz}o(T6Te5XPsed*$K8b@t)b*?1SC0?+rWlyhY!q##K>@_|9NNovZP%b2ScjuExU7)fhN( z(^5Ie)=m75gV)W;^zEd6C&A9uiLi5Z0_4LYTFSJ;EHDyM)hYM%=<54=KZHI^ZrAaYw7R8TuXlw=5^;+Vat13JL%~B->K=} zjwzSvQZCb`TsjCd_V&W8TRUOq(pH$cv=L_Pt%VtTD`AVZg~v#4dgo?y&Q3%9|Ls`` z&Y^`Q%RCnpZh+nb!W^^th1sj~33CkR6=tu_Bg`?KTbT1`F5y1#H>WV`G>0(nO=cHn z-DVT!Je*aSb(}?*Ju$N|>pGJ#=f;e}tmO>CR_ECf<52(q{H(;Pl5K$8D#EPY%EGMO zO2XYBzoIZ}xq>ijxx6rIxtuV2Y*}H}b{S#Tc4=YOb}3=jc1dB@wv#YxyM!=nySOlG zyO=O*xu`H}xri{=orQ(1wzGH**oTec8c>^V!mMRiVb-#XFl*N!%-VGpX6-f=qU}5&`AYu04Kw-WY*;tr$-bk2rp4@)wIu8)fI`khZpPTQe1G zjTP76uiNQj>C*NX@vQr3Vb*<=F#CObVUEK{Vb*zsFzYy6m~|W`%sLJgW*vtJvwyb} zW?i=xX71Yvvwyc1wz|$?CmpQo0pc3|d%5Y~{e@YJ{e)SI%Tz1YVqfvB#XiEE`+E!b zfc?FMS*wY{9LGI{*|U2Hb1m3im_53iFl)Q3Fl#$On6=$Sn6=$mn6=$Wn6=$em_0jQ zm}|uj!dzF!30ti@BIlw0|Ms_h*7hjLvQ|e5vo@Ej*6gDr#IqKM3p4k_gqi!H!p!{; zVdj3YFmpdhn7JM(Y`Mfe8P~E?#4(qXh1mxu2{V@yg*k2~2s8HMg&FH{!i@D;Va9rl zFk?Mh*kWyGCmoFKY;o-2vxM1mX9}}!X9&~R>B8)Z(}X$SP8BxaE$yU(^Wz+G?0EEIL|Ms^!_LzOn(Ezzuq|0%6S(xMUk}$``J_lhgFN$X_F9>rSo)_jg zJSWUtpB3hs_KYyc;b~#k;wfR);z?oF;t64ngME&}`S-YZtJUnN!BGGIyubGg$+C8z z3pYUTGhx>5Q(@Nb6JgfwV`0|rBVqRRhr*n{9|*Iy?+des-xFpn-xX#p-w|f*-WF!< z-V$c*-V|m}zah+8zAnr)`!(SSsK=|qR@=Go+ED-h{I{X#@bSN1jg3vTWs_AYYX|=P76!66YxU9jD11j2Jj0AcLAPXn0+#zFmsz% zn7Pd(%s!c0n0+#rF#BXqVXkj;2(xcy7iQngCd|H>RhWG4e!g(+ab1ItcT9PkUh-s}^2kj>RI*m3nL$aU7?mg<037gju&G zg;}Rg!mQI0!mQKc!kkZw2{W%ng)Qz@9ye>WytwWCbvrF5JQ{e}q&o!MO5%0|UQu`t z;1yiAwZF!(Tut0~$gY}fts>0zbY-{I(shPHcMWm$ySmFx>oSaaeR0fpJ>gNXb)V)V z^Iuo|Sn%rzj{{y?cn9EnrO%qIC4Lw1cZ=T>cuny(HXZ#n&i9_;7<&)lVbFa@wuS?D z7e5lXQJ6LECd}Bn3hxB@F2anpL6~FES$J2-Zz#;VZy;>Bv~igpknJmub?+n0*m?`^ z4*5rv8*TS;+wELt7-R>D8wEU2m}9uHFvn~oVdgzRn0fXWW}f|oEiZ-Cba2giQZI50 zwv;UM+(MZBv$-&1*-V&q*i@Kfw23fd87yqE#GK-_WLt6cy^S#Ivb8YBcPqCQ^Mq^E zC~=%)FQh!TPqs%2Gp`ZCoEO7|S@U7Stocx3<~c-oxBtrB^zU}UmUjz(jd_iBu9b6~ zGvmb3?^t2l9pk#uv&?yCam-;SVUEp?!pvd3Fm3N3Y%x#gadr2zo%RvOIQJGF27WJL zj_pKY&h0&gIo5j!kAeQ~!tA@?usu^l38 zvBaD)6ys>`^Sjy$ew-{>#(I))$ePHB!i@CFDy`sp;SI#q|b0Pj~?Exx%c&Il^2!&K71Y zX9@FKaHg>NR=Jywfqu5r#o}oDB4Mr-lZ9Es3x&BpTp&CV_9h9lhO~9LIL7^%<^gS8 zCZ6MQsqkK~eTmy{?XPkDy;>aCqN{{Cepd?9*A;FnYSImUZV=ZKcv;~-z}E}2Pp=bZ zO|KOm1o>-(ttPGfwMK9^i{seeB+RkDG3j=O?rq}e_f}!vuixUj9bJdI4~b*_9u)2Z z+YbnL1-@UH>+yZUJ;2{9%$nUJ+!y@a!p!R~;f=xHDa>5%5FQNvc44dG^!^%ic}83p z$UZI1*q;(+?6<4$So0^vb9|o=X6%m(bG%m-W-hA|cHVau(h z=gw=Ojx&_A_y;%csH|<4=UUK>lN4*6AZ*_U4DetlI~|9P{^uIp*&P z_k#U*g<0o!gjwgeg<0pfgjwe|h5N(a8^Wyn>%y%2Yr?GitHPXzuLyG-UKVC;UlQIF z{$3R3IK3duweWdi8@E|qhqeDr9Bcopa2MGAMVPh!S(t0kPr@9FAB8yY>l`w1hr7&yxg)nRRxvXTFl(`f zF#BkAVdlP?Fmqp3n7OVZ%v@I%X09sC6>e@04iLw_ z=`YOM+53LhuAlg>s9j%S)~b&%Yt>trebY;rbFQZ_`>2O7YuR0xwQLk-A9WLEZMzDy zwq1l-+Xi9QwzDv6yP+^^yMZukyS^|#BVJF~YTLnU(iOGYOdM;osW59XTJx8&Zz7(t z4;E(ZgM=CDKw+-q8w)e`jf5?>84=SE|Ns151tTQOaUU+saUUk!6?TRSbDW0=Gq>%8 zId0nubKJHOX6{=H^P07lFl(`;Fl(`eFl#YJwPfy_i?`f6dd$paXL0PgorKwQI|?(_ z@xr{<*g=@Fj}vCBV}%*(7-9C@Xko@aN|@ugy)bhbDQvmS;+PrR@k7)%%>5Y2Tkdlphv@12yv&u7Wi74{W-Trk?g~4X33FVJmv7d_ zK1X9uUn+U_^d-Wq-NnMZW?dxAT22;bEiV*iEiVw}`Y}nEwLM>$wLMRmwLMpuwLM3e zwLDvxbL1>x&XF^PS!-VJ?-j?|-XqM~-Ywh}wYAR;S=+nBv$l5%v$l5#v$nSjv$nSh zbG^7#n6-UReZ<<{BA&IpS(vrFNtm^~QJA&7L726?UYKj~b;7LewZg3JHNvdz)xx`= zK354_t!MMsIHz9_$J#$H%-TOE%-TOI+!eN;5oWEQ7G|xV5@xNR6lSfT5N55<(^#_B zkBeum9}{LDKPt>xKO)R)_QS%g{X@d6{e!}+{R6_R?ft^6?R~;l%UKJx{8$`o`H?Vd z`Jpgt`GIg()bf2{*77}J*799p*76-;*79v(*6uA~*6vMV*6s~q_U-G!tmSLMtlg`^ ztlcZZ?9-Qp*{3fFvz9LkTkWRvoSEA<;+Wgl!t9H$gqhoy!p!XpVdnO^Fmw4#n7P>J z-i-BA&1uH^spL8Cp9nK{j`Q#0IL^NbbDVz_?&>+)=@((HS3e7L+5fvyc83=6kchg!w-BPvHru-ygy@RxSLs9l*^d-^O+H*Vuy#iQ_meDBKmg3kY`u zo?n=CpHG-|o>!Q4o=2EPkK zG(|s_Ox*TC=)9D35t~uNoL^ep?*9LFTBaZyJ+xdCH|1~5%epMR z^R)`xLaw)VfgA7KIt4ENuDfo5i<+)i;3CfT3*5f$Z-W9CdK(tFm}8v_T#Q9SfeU|K z3S9JF*8&$dyA`%ATRRtH7QJdR+lr%|Z7O!QD{y<7 z5gkJcT=*Ya;f58susOWKji_)V3tYsqeT5rU;3Ae%-vu|ipcnqe6u8hEo48oZt)Al& zXRn3L9SYnYCZJ<{feXDI3tZ^!RNz8y=ftspc1fJY5H|A|#XX0~Oh|e$zm3~9ah!9z zCC=o--|iJ|kHj&DJrhTp6D!${r&PF8E8J-XF6w=HfeU|UBre`RSxjdpj&YXO z=J0n`K`*$o6UW-0lQ^yy=O&Kp_IZi3IHUK^PaJJdN*wj_eGU8mf}}Uf6#cj`arSb= zGTFJ-Aq?(PalDVaM3`&K#lq~bi`*9HbNIO|aqPXz6UX?jNF3MJD-*|gc~#R7~NXo?C;xx&4YxXw{bonNhRuO-gZqF)9lF4ho} zc|CEwp1zSdlaIc5v%tkV@>b$_?Rz_MwE0ejdpB`h&)-WNZN6XOK1f{L)0qB;iKESr zD%{5v?vo1lY2tY8{48;de@X79I(o!=|m9~JJ;#Bn_TN}R(uW+Sx zDX#YmRP^%u;06oSj|G!nw;*(0$T_aHk-z=s9zUyDIN9O#YLUd5i+HcFXoXwMIrafymvrwX@Zgu%ss!?>x?`%OnQ8VsdM6Jv!TM>lyacIE=i9$bWI%nbxRySlWa_!>BaTC zdxh&!;Y#oS!)9sRgX>wb*{j0!u5f)4*T)j{qi^Eu<@h^AzY5pC!sWGywH}c4;yt|i z-6(O4#jZ)bujb$8Hm>LmOkB(bvo2I1i+RXET zT{(?KK@Z45hXE(3tmEIeK%`GZ=rTfm%+p?m!RfQ|P z&kdVfCq2HW*d}r8mu(Zr@!KwOaW7}J9+J3t-o?0~iL=+@Gpb<~Zg_2Ft|w`WCf^$NFpWj;)-=j%9f6USN{l{nfg-IImCqm$n9@OMn&X!FdrhNixz$Nkl@1-;;oOB}Cj$0v?|CpaN- zysn*?IEyU&m98_vomA0_&m~%zF+W1@^l8BFBt#C3z-If-iocW&a0ifhz)74G~B zH>tv1kT_mPFH9U`DZTfOIOBZ**OcG}X)WU4l_ytxTvXxasBjl2j`vKLB#t> zWr?F+zUFccTwc+;qQd3R6U^U8`$xzAh4nnRQXd3Yx()=F??KH^a91v2XzFz_xIF)< zXV(wv@qX&6#8IzS&#PUJy%n)sD~_LkT~qNy_fZNgU_I zjmZZ;Grc)+?Co0;$C}>iT*?3K;&{!vE!pX5hWxlAarW{E=kod#?`w^|v*17U?h-fD zZP@AVijR8|$Mx#o#F?$|abMziy}3Ve)O(=9{acua;ji=_GPr3fHXp2T4<(Mh{&3>h z7mp;4`Q__Axkr;8-^)LiIO;v_9D6Nte?lB<|769^Qx)#%3inLn*dxy-j%)REiDOQ^ zD%|rG?u82XV&WL*ONnDFd7a_<@Nz}(mBhuEnC({+M}MXFRiXP@(u>zk|MkStX8zop z+#5-cWAtX?sP|Tdd%ME*sc^P7^BP8f?^L+b`}e5LyA{3sJWF45=*N3WkMC>VPh8w@ znd&<2?cXWaZ_eWnlAcM%x{~iR*q8IY_kjJQ zqg{cESQbsWn|yHXD|+_ayRAK3YeTPt^g0{u$FvnY(^a_qcSSQ7K03Ny$^Z1?_&M1O z6(6N*aQN7wQkxkoHb;lof2oRB9?d_siphpXB%s*z)I|65y&U;%Sm8Q5*D?%2&;B09^{^q?A=f2wtYO#0F{f^c zd@$91P?vO^#3eZ2@f!ua1^d2r^&Ck=(-)0q?n^(B}`6J`pBI)tIVavo(uXJx0Ic$~mcsenWET>rCw=Zz9HjPRgpCcWeII|gi{(On|5MwHOV=J6}Uk2}ABSw35Tt#n(#C0(R zKgK7{UXC2%{d5bL9q4B}?O4&vd&5mWZ067TXmh7zvk98Aj-3-{uZ7Lh-*SUHDAkO5 zhZK4;^z8jQYfZiU`5?JnT#svP*tC7eylza8>?Wq{N3Ab*t;DligcIfR{$)R*D2)*?y`ORJ7CMJKKW6&*^ z#PtNXcj7h$w@>1Xia1Mu8wzf%N}TyVp6lzr$sf7>D%}1R?tlt+V1+xV!W~@U4yka@ zH`R8^=L5{|g$j3Q#pYqowQbU$`ZLHQ#r5#l?X+lzslPk+2=TqZA1=H(@O+Zr0(f3w zu300%9p!PvT5eR+OFAE&INm25lQ>?-j!hi>9hW%s7oQ{8du+x@?)aqFV5A=>B+gzA zo3R%?@{89^=0w-y{YCI6i{o58sUmCd51VK+s^b)w4WIT}*eyK|5nPK({-;*_+p`k% ze_BP>zPp9j`O_=0#x=QApEIRLpJ!C;oK@jER=Co&Hfoi>dx$R3zrsyQ9OujhiDRuVtZqrXckTz*D@Ywx8My~~^{&5t?7u@;vnJDXW>eq50_dpX7>e;)|f;459P zbr^!4-K+3Q+m%Gr=0iD_dns4g5OiPJMiwpT>l=Iub$v{6VH3k$HeyrzpHrm;iKZY z4m~2=5AwT6ek0(A#Sa1=FP`s(9}>^?;6dT$?irt`a8D+V_oz=LZV2LeI&ntDcOE~J zIIg$PCXVaDbBUwP=PO+4`-&o#ymy-ZUPv~X-;0T3EH70!``bevhlr(gUmM)XNsnXy za>ZZ#ZL)>A^&@PyahqHxgKH~}wRpvSge>REYvTHcah+eS#PfQEd!xdY?l~jYLsM-m zl;Ga1*j&28y;b4tb107C1cQFOUD10dah%`pCXRFYJ?EnK<|^|0KpgMw-cNQoh94%5 zb_OJlIenBk>U~_{R;Y~2CrPi-jQa6u;_T(f_p=ID`kX8DN}pu~_qp3F)%QzrT&KUN z_%D6#6+Yr`za<~`nHcN+RmDf?UNL-}nd;5!$=Asy{e6=-*85xMN^#qLWKSP6JAGHN zQ+hTg;0XvvdLcf$+?J|dcTO{ z8u@dw!+XL@nhcchnLab4c;<$BzgBF1?0WH@hq3%AJ>J{=A~D;;1*La}jHI#5%V)<~3Kv&OC`@t>#S}eax4*P2qq3#L?yg&PCkh781v` zdBG&hwQ&(~yuVr4WuwNFT{LlgKd@NhXmjzzG1es#$KL4VT&#`MTSi=?z2L{v!W{FZ zgjvfa-Bznm&^dZ}d2yZb`f|eT!)0AJYGR_HyMl8u4!q8;B#w1hG0D=-D&lyJTsg@y z*VV-FzH?QVE#-uU2iqn)?DZjuV_riO$MG7L zIQk!+IDWo1B5@q6kF(z^Bt+CF9PwI^q$2#sHJOMF{ ztJvMK!tIneuJt=R*WUebJSK``o%a;x+}%T%{j|F<=iqL_?CD*FS^o*b9OGSt?H;}@ za^A!Lf4J5w#=W0p>1*F)dmmwr*WSXcGX)bIN;E*6RS-HjE5(VK8|oM;^lleMjY3i zqlLNVA0@G0s>`-BAwQxf_I=hIr?Vx?>&#ig?Ef=`IX-6y zb3U9dylFVm`7~kX%eC!XajgG2N!BzYPWz1t&iV5yvWqEy&cX9twv8F|BkbBU8O(pO zWVzN{D9n0XAk2Q7wbi9G6L!V|b}B{a)g>qR*MbmEsu3 z6~Y|D%U!qB4;?y9U61P}%Q&tRW__*|=9pe1%s8(W=A64q*!rO~7WUjJu;~f+Os*#|4x^UaiiUP z#Bna&EzGs=uH=WcxlbJPxz}Y&ep;&*Tz4LjEMvSs`C%_TB#!g@L6?p3WFI~*j(#2! zW`2(fGv-HxxgI<$Y`s#-!JZ4|I`pJuIVMlIZp6i&eNG&8pA}}EpAlwlpB84{Jtf>+ zT!*N(ticOO_j%VX_2_2OWn3>ymi2l`m_7Pp#g^@JX3Vcime-b7g}Fw);(U<_RRZJ$#R|gM3}YzIN4&ZUy7sOFNB%H=fbr8ncIq- zXzLqstlQVZyrzETx-pjY`<*z>iEmvt^5^{eNgQkUqcHpU2Vu_h@7-36AA9I$=OQ2W zj6FZdUi?+EycYeE>~f9!Epe>r@6JUXSffA1abEqAWLdMn#WBynTsHD#js8g-`{&=} zgZ^7AHFYmdBW!C^_~4vsnKhsy|a?OWP)+nACcF|I2m-4%sdgB66C_wvHLwk{`Z^RBc8?W+Fb8nmipIWDUR zvu{>*TahDkT05=b!!Gk`Om@2obMAFb zcA01Q#Bt5*kvNV~Pv_ctwR!F8FOGKl2~)SPFt0Ivgc(zB;Z5DHoq7rL+F>ZhVkxch z#zkGoZ7e<3X(M5d)d06!ns;-^F8g~E$#ya0ehe07j)R2Rmjm5)^cLs)rq0D&=oUtG z-aOgeEXkT^;1m zN8PbWmOVCJ9M}3CTsGokukV=bu(msk(Qakm1;4q;^Ou9 zaLKaHhq-R7W#%PvI#L{SI>Ke6Zj9lm#IY}quJ~A9@v% zUDo>~ajf@=!kqsnxNhXe_)m5&YQ%9mRUGHfDaj7G(~_N@R%}1c5Xan3Pj*Z!#^ub! zaeU9JaAzluHqUV`>P~KQr>Sdoo@7~rb6q#;LtXnGK7KBBeu{zjPm@v%TpuQjifStB+EX$-F2hpycfJn9P4wZ%SO$4 z9lE>1-IF-h>E6V3iGt~TpK}o#^Rn-oCOcerA5M12Jt~fU z_DHhBK6=c#=p*L*lsNXylfoS9Cxp4?JnptiHF`!|Pt@UQmo52yK^*-)FU-0;C(JmW zbz9+=ee+V1ebHs3Z+Ir*p0}M`jeqb#;bhcY0y=RY$iKFoSwemi(S7*)CB8ooAIS zuRpU0GryVLR@8}W)NG048a}&orLnQSeC&rgCClr^92H&L>&IT7Te6%}b0uBQq4~t| zdO5E!$9*2xE!B7daote2`CT^Z!Wu5Mbdbb?PL{wRMSPm)uf`--8ymgBp!IQGU$E*rIEUaKUI^Ji7(O1)rvO6hY=$+DJf2=hKSVP3Zk#hl6ebGa6FcfC?>b89SE!yd^8 z`=w{%=zrd(g|A#c`z{&g-7DE)FZ51agE{e|PvSTS`Z^bVVJcCZ0pd6g{e`*C^>f{* z4}ESdj`3~evTY+M__2M+TsJnAEXQ{fVa}1k!W`d0!n|G#6t?j$?zDT?%}p$3@pTm+kP8^w)d4R z52c&3VWxzDy}Ky}_QJ{3WoNQg_y3y;Ls~0&JwV<6##IY6^C)q~N!%mlpWBivUS=zlq z9OvfcNtXR`r8w$#Skk;rdG3d~T;+BmPsVwTWXNUFTf%2G_@%#4(Q> zg}Ykle%v6;b^3a@)y8@D)$PgFZNiN4R$;COw+OSRZx&|FmUka^x~pm6r___&9TmMh zE8JZb?(PbAPlda;!rfQl?yqnURJaE#+(Q-a;R^Rig?qHZJyzi!uW(OPxF;*zQx)#% z3iph2b9#TVhd&g@`SF1;=g0fPoF7A`oBBHYo_O~ByTY75?+9}(cw3lr;w|A`==(Q? zxdyx;%ys5IBWS);!HBG9Umu-*YQuBi)#b*J`<-~x&I%ZCOee_$?4BR(~8w>8+#Ek~`UE+p-`#y2If%_qGVqJUEvO1!kkaZ zk8%D}(OV_mV^A-Ci@Alp=ttDLO{(XF5YYKR&DZ8)`L`0w9e$r%@)zH45ggakX_ooV zU-XZmh|_*Enrl@Hmkm3V{XEr>Ygo&QUaJacd*H1WrPvmgKd!^AD|XsASF&T@W6;~& z_|djvC%)I9l>dflOwC->F@D3k*b#r;- z@1LdKbg54^G18BjD*h%{=7R0x_U$%>;67(@*;bBu?=rhM+MP|9_ad{p zZfQPr_dH86*l#y5hB@3$N%mpo-8l-R^PCl*b5*z(X%~+)Z>`V?|P*%dsBjpdx46N_$`Z) z9sAB?riq9=l{=BbWJn3-`FOfLrk@rrK>r{z*W7V5EFPUT;&5$2Ui{l!&R3#qU zqte_M4@+arzFDSXv-CS8(KpMcxXt_jIgK&>(hw9adb<2taG%a&?sze~^iw{<1U-dm@lYtP|x4A++|>$jfk zwm@u=&ll-BVIslxQcRpDjAw&nlX-8LINI!7;TjUxXo`MxNu0eLadves#(=}`mudQ)_xa(c1Km>8iT8H?hL~ zW2xiwYv!<5MQ`s4w@-!Jx5DjL;r6d^2UNH=O$|2X^Lggiw!$6gHfL$7&(!x2>vx>` z?~LmSbIm+fy1d6YMwn~S(ZY>x*G@+XcL!ce_V^z7Nby`hj}Yz!`NM^K10N>LYtd@5 z$M<)u3iGqNLnY6iKSY?HGps22O<@0E@%+2sLBi~z5s*F3xzf60zlX;4;`n68%*1-} zWMxe_A?b0Bo|rhZ8GqNxf7=*hq#q|GyKUK{7aKlRSlx*GOP=d6ky`)*|$@36!1K1X`Ab9Ti}=QO4qw{ugh z<}vJc5XZhbFZp0z=U435-!EuqQnJIbyCB&iH(A`S7N8#&R%9<0H#P)yzQ|>Hj~#j1 z_r%iArIKZhE^*x$2j+dbIO<;JvMt;R$vv7l6N&mfmN;H#9#7n0 zw{52<5*PN2doppX_fyVA{w5!~&uAWwvKRb#x+43$?wJ|uGs!1w{%qpReE78Qa^Sq- zcs!T%IKQ7y9BuxR?&rzH?`}u_%=bmvnS3zTw-Pte4EgbP;%M`o z#FR=Dpf-1imkhYI&&h5I+HMZGOyKYmJjtoP4}GnJ^}FV2z`_=7~#^0U^>jK%|CEFADw@M85j97Qb{!x*=SM_Ot?4OmGHk4h; z{w>)?=>FxpF-MJxnmwp|nwtmtH(odn|H|edQ}yGYiVyo8m*#xr=g&A_rdf7sJ#sCa zD~*voW5%&=C0TwR($aOKPdJ`!#Bn@ZC)uWD!*gyYj(yYCW!pM$r>L_%cgC1HN|s|b zoiIOBnpU_Gb~*^tS9@WrY2-lLGl*kO)4Oat*o(OAnYSKzeKyJRZ)CFybKGYU=6W}? zF#BL8VUE*`!i>{U__p`IjB|FEZEr$;gslDU33IgfzuiMf=Xs^eb#NYG<}tT0>o=D$ z^O#eZdCVbfakqhwh{c}oV=N0wmU%27+y{E|3p2L)gqyQ%zq8T}x{FGdLbj$^cH4adI_^1>~k&F>0)7yaZky!ZaswAo85(3$3|hkU+N}ob#39f zvo0GY&P1XYC#mM_X>tRl$C_>Ic3PQ`4%TZEaozrVx#{1*!i;~A>qdO6{icay-8V}d z*OASgi+Id&&Iv4S=Pe+O4^=A8IhxZwy6UW%bB#yC-t#EHNHQ1DU z6UMfK^f=z*+)mVvwc1G>b$1kIUgKRiYQXp&(Ol%*+u3EK&dhU{#PM%o6P%0q=)<06 zW-WJ;sfm(put5CSv!ZLyK6AeBEm^K%d%14Zg)!}u zIM#mOickBjoIdxLEaTrV*=1c0sMxj7**T^MNtW>*=(dRg~ zP8_dg$0j?h#g3^KKw>ZXrm&-;USte04`%RJN?wjp5L^vnzNp=|Dy@{ib z5t@t5`B+x(xmknz+)mVhb{-JNF}=UyGxnI2a-QTlmpJ?VjOJn;nbwpClYi#%P~s+- zSw9|5oV^_5@koWU{Uq!`_SK`V7jnc`W&$>=TJ&Y__+9{+ZL06}_ht$NBzr z;@G#(B#t>e>s;i^IIq??@tX2nvcqwDzQVnbIQrN<^$_EJu@c+M;^_aS(%whZcFYEn!#ee>Oclvn4^-5zpyL_;2Z&rM~mAF3U z#E-WVXD`Rxd#A#^>s+fa1ifojH_pZPCCh8pd&w@x_ygx6XXf~kI9|6uOm^t=<79{2 zr{WqB_a{l#sHlPM>*4kHvx?8p6UUstNF05X_5g+d?v=l(eC0jZ0=a~qeXn|ReI0vg zLa!V2zpTW0ozI6D8&eB^ZM1KJWA#>#yr22T zbt9MVsKNK*dI5iz{IFg>iR%N|ABDM||KPf9AQLsSXDXV8$+_Po%f9?I+4@D8bLD5Z z9ph+b!nW=6vGI@m$mM-KzgO)1;asW4d&vjq)nAh3`21P1`*(%=r^4CZA!kkYWht`QR$%ttSkGC{NfB6^%$NAI7 zZI<#}T{cr~AN_-s>$F*+S#F@FsYr4cS_Kwa)U8y&NIM#0ZWQVf$ zo!rg!uIMX$7nD|*BBaIqg}cROv(l^-E%-|NhIHjiXk z=edP>FEy8NV;IqSPGNqQIft<2Q|dk21J(_?^SW%zKcgb9(cUv9JGRe@^J_l0Q<|&x z-Ojw0EGXH&h+zR?_VN5~t2EB@saBl#3rm*o*%zwl+8#6Z$fA;EEf#UzQZ3$+ZgVx= zC)r)B5`*nQWBf}EyObb+qqMX5E&NEcPGmNi>ZdaTh3 zZYN~vV`Xv7Yo&^;eaA5Eu2PY;eSb}J`~F#c|BmNIob0 z7fRE=YbDNJj4@fe!c9~>&CL<}-Z}bM$L*Brev@k1$4vRLuFIDC_yEbW2J1_fj4S zH?YDDs&IoV+$I(7;kH(UDQhU}GE-ASl(|hSHaDwqn^(9koNMFb!uytO#C7x6?X<5;qFmki=1MXyQz5SJxeuIQH=H#8GcV;`mwm z$iz`^nC3fkCuh%MliR_dAKNFJV}sCnRN{<^{PK4mvM!^O-tMqDCULZB&nYp^5zrf( z^oD^Om$=U0c1Rp;j!ztO*wMLITUnQ##qnCTQ?f&Dm&C=oZy`=d9D8Kf#F>25e)=@l zjBU51$9#8BT-Y<)dnAtYZqLN^H3dH=CeB_CyL&koHEakao%a#P_au8)?5vdXrH_4+ z4|2y!kJq{VDn9m49M`1-64w`b9GE!zKPYj$?>IPdJz?{Z#L?!V6>fzkO>*k%XpF^S z(qms9Ua?d9-rca1_pvaxBPup+&q;H!*&Y+t=18~0d!q1fYZ>Rs(URpoz)_VLj&Uyf zk+IwU5YGQ&E3(Il<8OH!U-4PmlM=bu_YLo4LHKb(#pj)pWqc=Ae4bRX^N#HBI}1*( z$etpucLb;NsTH3O_#ADnw&P*vwB(cPz!~Cr%{#r~Gxq6}e7>tZ2f^o=6+53-Vz>Q1 zePQP;w^Q3gOlmi%4P-A&Ptiy~=a_6IeWx9!O>uH@79vC!uY^13Q}+hm9LZMO>Z^O;-RPx~;WbJ+S+wwhcycduk=`yS!G(7RihbNMb| z`npq?pSRp0%xlQ)!mJ&A-Iut&@OyvajEWq#@|wq3@xI~#*DLkPgW|Y0K2-7XaK(;& zcR1(HBNaPsD?Mj>9#|XZ_-Mtyea|?*Z}PE<|Hmub6BQq}M}W0@GTGr6Kb1J zF30@c=F~ge=R=G!xp%y7r7>$NfwKO&aFsy%(YAX(nyKA(JYp1)YJYtMOe4S89zoDVNmbnSc5 zc`bQWvV$UUonNWwZXsQ+d#_8DF~8=zrJi}ixl+$uB)eRv->ldf>wQ_W^Nj4U-fy{` z(%9HD-+htyJCbGJzFmo-i|TLvP>P|G?t@s9cim1&b~eS!zqP&RvZZT~J*Ukv{!p@f zAMt@O^mFoDo>d0RFOtOsQQ`ar^#OD=mORsTB@14rp z6yI}SiglQ5_QaTfDgP6IzwmfUWA(Z6WL{rYWIs@R9GkCQwluExOf~EGtz>&6rf-Dl z_X)+pwTN@~JNFYcBlmsc%u1|NKUBD7z1F3g*mKjY-;WhLliW_k*4F~?<0scE`LJiK zxi=={5N4cKRc6_F9ZhJjY$C z&1za}826tQJ8M+lpZ!(QJF_wm*Gn~I-0^z?G2Y~Qx?X9H{;gQp@Ba!j$A3K5Qa$XM zbJl#CrayRt20+yO_$DJ|z5M$C%&SGkpM75#`;B@lRJg4wwX*LHW4^8Azq=Xs zqh*SRb!?qDQ;GcBB#zgHwuz%&y9(DnaaMRUD~SMpr(1^U|s@Un_YgfsR{q_R;dhUvy(sSR@1M%E<$;SZs=lIN1 z@iAwU;FPt4InGgN^7z`TTa;pCChoR zfa{j-Z|qrf#;}NFc@MU*Fne|(w^h1soRaQ+Sff)boIOk4T>ZYUaPgdYDHeN9yg8eD zIv2gfd;G;x{)<-fU7#}Vi&ymG`SI{?kw#D0bKv}Yz!Gk!H11ofe>gTvN|x8=PAP^# z=E#qw636!wODE1$BEDrR+zfv0kMU>-XF4zIdZnE0d1%JHT*XJ}dFY639UqTUt`Q($bD<;m=!e@TYn%vgPtGOOLv%;00r4E0;ROZ7<9%mbK>qqF>GuTaS z@2{^WS*{tY3ik>nomUYa4ZO13j@eQ7Dt7JJaOPqtYGBWTH^Fee;W8&~0WsBq(*o85bob7Eg{ zT+8>?zFmP{(LwCQ# z85Qf4{eBMh$bC^+`%cik0`9%-%30wUNYkhG26USaYCfOW@ z*pE$|QDO7A3fC%)H-8J|_@u}0mpmbH^ml5Z_ac^uQhpro6O&E$&`A~UWar|(i2Ze{ zINCWS*;ikAc22A4y`J*rn4O+tqrcwL8*eZ8aYn_*nTZ<_N;;pFIM(#+#F<)* zMg09LdUQN&o|E)gi*pm#1YP&M#O(x|=O>PHZ&Kn$LGOaZ85KEP=v<6Fxyj;Ki;I$- z9pK~QiXD63&mOv@V&~Gtjf9WO631AV_xs8i3$AsSCq2e;MdFz6mCi+NsCTtE&gZL= z9maM|;<_QHYZEsR+;xeg&Fd@N4HfRj3U^cDcpbetag1ewB@IpaoP^isTaq5RTb+yg z(x&Y#;C;bu$qr@h_ciFBTxZ%;mX-*so=dO+_k=Sq3C z(VE5U-rdPgKiIj)xl-)*I~|PmKFP9o?oD>hY|Q!3(>04+?5T)(M6Pe8&bHse#vpQT zFzCkviiNeizY^PniR1V@lsGdVIXztA?n~E9>cww=MBLoqPcln`rVVTY5N@5Yleb*R_LbEi}T+bc0FjOhHj+X-3L;7xISmiG;pExlK;-(TT9;5(A-2S0BMa}9mVZG~SG zjXE8sIx*LGUAE+RL-jf5&j*rajo%k$T<^K97<2j^;dLvGhm9*^`bc&nUbFF`+b!ku zk^-dc$CVi3H)KjN*d7qBE1#wqK5@GxUE34V#RB!?bIEc&`>YaE?Kfw>NP1jzzDykJ z_f_I7tf*1!9Vz9Jf2-yj*t|5cbh^owwF>-nRK+qgD%@*~E` z_PKDZ|B)Tm>u+J!=PzMCBmbxSEse4L&P~^FXy+BC&hcN@EzRNBV^bP``yHBYu-ig* zEx%F>c|T4o$@2Qy(sfII&QxtUzO7w0;$_|1RP5N^9nP7y71`EFwp~Tm_V#dW+E-+4 zuMhjagUd!PR>A0rgEVf<^~+>&4MzJhz3k9#M`8Bmbi!fRRHk*mrLj6hF>-#*C|O>| zW^mooSlJ#V)_5kDZ5x5=9C6t95^(Hhmn?IdO_*`aD$F|1A{@5N{LI45&Akg0S96^0 zV=}Ik&p7Ed(I%o}4);-#wLMOhowE{;?Rjc0p1IOJA?L|l6`O~-&C-~5RA8*-JhH>u z&0WdM_9Jz-aQ&FKB0GJm!F-i?AN0>FO0~58NsMp)icj0G#M~CB_-x~S&?b`6IsDju zC63=BvO~KI3v-SwBs>sy7IeRD@mkon{Z3rNI!U%aHfj?S#c~@k}T`BqHy@PSXOXb zrQX^}wz%%BT=8prxG1}tWZ46&3J)}geymclWxokS-8CiKHxzYVLzvHwt?ssB+!)7N z&Xw}By=Cmytt30p z1pL_2bz>~-<TP@Sxb|$F;$!UFrFga#=6%yPu3H+brR0Zo8sf60 z{O8dK*M1{~H67uyr5f5kz~*W=Offab#qXhnO^(^fioeqS zz!=Z%D|)eau+-^ECD_FM}lRARH=Z|P~O ze(YNDVfzs2!%+C2)KskUdw%OZ)^d0G2tBj0Tg9jCHEhmj>@6&fwe1mP4fb#!rLnfX ze@&^GD;;}Dw(-CJZTfejFt0^>x~-C5+bhU*VDF0TQ|d?7ejk^OzA)!er`T6m@@e}D z>2u$ToySt23`Cy$RrE@G5F@tOgIJ2q_8IbCWB*D$Y+qr2#B+elmev8=SI99sNU{-+ znK;mOqn2G^_u#}C6}t92HP&yXTw3V!G1mEzik(9%+(wn2?wrPlu^m>ixm8mGm+y1z z_idV+LpS?1qm=K@vdMXSxce`SN4r#;BPGjq?TAY34|xnRR?N9;r5;Bmo4ltvI&q!N zs2|5TSIXIbi--5h$0j@Uaa`ieeDsm+nQX4+*H?PX_7<{NkFWT!y@k#BIH969r^c$e z9=H91tksEbr!?laU$DEG^5bO5a?DSv#L%-cuC~{(xx9M0UMU{icgXizr^-J+J2}PU zDUJJJ)slIiR*60KFh;*}efy>od)~8XSw_E1tc>L}`rCVR^?0G;@5KstdL_Tu2U)sS z+CE6G31`VpHw({?Glh9?cSiEV8lUZ4X*~W=J?Q70ik)*4XQ4!0j`!zAH>&^uq4_nPOq-%>By9!lmqS+czMys)Bcdn~DYv1GX>T;#f?y4fB~=5UG2 zMz2~rQJd1<%HZ~@^yA{bew6yd_CzxF%jLfdYIvD2*Vjucx!7;;u>Mz6WWQ9iFz+i} zwlp@jHtpQaERDVG=cL`IWruZqN;q;b4^O(E(iqv^ zKi1=U$uj5XgxP1$3Ny!N+;*w2Zb-IYlq{bydBJr{eP#O&SsUi_Qbq58%00)vZnG4t z?K9*Yenmcb&3xH?mUM03A$#mK$+AygP4V)2`g-Cl>&P|sD27kkd`o&I2D!4~yk}*Dd*3Q+|4zkRKmbWNp7=mk`wXBbSZ&!2bI{Z!$8u5iCpxL+&WZx!zM3in5a`?JFRRpI`waQ{@ef1R7>zm}Zp zx;5vDQ~zCK4q^8C?802%XA|Z;nN_&a?b>M;;qJgQ3-g|CCShI|XB6hTK7%mVf$4?0 zu67jW8Z(_R*V*l6n)=>iTJZxAPX}TC)`pJP`72g*n&T2y=aDEzEn9R>Hhi zw-n}Q5}V94^|Qb&#K(Hny_5sdBiPbPW+g+;^X7WJ;Hni zyh~2yXGjbq2RY;$pl_q*LO?L2t>#85Oa` ze#D4-DD;+=9{1fXmF$d$>@tZP25#BJaU7ScaLXrdH`rXkxrm$hl`DzknzUlFvj^;~ zoH$;SS4o_yMcwlE`)m)Jt5)<@t8lAVxPyI8MZS#1_JncGTqD^Th}hOl9M}G}oGbaT zy~gafwJWl=cZxpONj@2S?4^piKrZ$x1;^{sx)q!DTQ6K=*bD1b^lVQN^>|%czoPeU z@@H0}Phw9{siw9kh}W78l8=qekRKZ+&Rz~5wzr3Ivre5WdJT!=X9`^s$3E#=;kr3j z%6F{#g7dU7+39a-_|ZLa_HyLgqvFH%3NhB66+0{YoN5zJ;bT#8w9{91x3<$NuQedT->iF`vv`_$mF~OK`SF z$@;4l-yvzP50nqC6B}22+J2?x>~5&Qxuy-4odG89#~`;`iphSThIwx)S>8i#;<}}= zvAs<6yIDon_BIWOfOOtG$#RUh6i2&TxNIBfN4;$?6XV}jvJH^iMwt0*EzI?HE4SSm zI$`_1RGXoaWuFfb9s|AY+*YY)Y+qG#<7fM(y2H+J*c~cwcpmE z?8r*I56Ldq$r13eeMK+!Rh430Og0C=&ZtVP=AYv_x*}`)s5mZTT(;Ejwoi)hvBpW3 z{V>*bOJivJqO2xukVEu}?Nwsm?Ib&lX-8r9_;_KiRXbGt+FqsRddK!6(atWi!~K*y zyWLXEwpWQY-d(bcc{gF!YFA;-%L#6~R4d!t#Qrf9wX(fS&DGKNBQd4}W|}&tiSo&7 z_MXBVzdbyr=)cb9%8$L1>|QQgir@AxS&U)FP{eQhmFRPyik;Z6RPr&|YgAf$4v~-U z=G2e<<%6FE?3Z%c*JCKfYkP_s!;sDgRAg-r&c={E&}B>U=RG;57B#Zp9BOX-V^2=l zVo*ldqLb3S4ZK*@(a=ij7i>N*a0AEiFGXZ0z2lw>>q zr`GiEk-`&zkEq0(e{09oq95bA_)@%sH3qy-JUaQ{ns`j&7~8RlqyHz;n6T%MmmWW- zJFa5q_R4+4MQLnLlpV%*LbBTfbv-F@)H~U^s2#`WRB^mUoRaLY2KkwHuH&aw^x}E= znA@~z&%-y@qnlLbmObCzoSr?;Zd@BX`4N4Te|zRE#m(#QnZg4@N#`?!SE3ZTwSqwO@;d;)y$#} zf7d2G_TzPlV=UJvjyC72)Wx39ZLVfFxXn_Zt}eX>Q}yE}$#w~#^Nkf&op(k88{wiR0Y6GjZl2a=I@0BUgHEH}p2I^lj<6-OwvNn;YD>T5Fg)W3gvuo8#;^ zO31M;cX{rmaj|FT*r#__>SWK-(TAak+n$YUj*DmCqGzZV&!CmK8{KB<9^0Nd;~cq9 zar5)ndxg1H+*7H6JwwL1dw)gNo+;z{`+&=q)`Za-6W&KXBw1dwA9USP?!7Bx9lsG! z(wncQ=Uf{b06${<^7CzvDK_T$s4#o#5n4J8JcUV&i+l=PR+sGioKD@g4h74|5T9wP(^e?_P8tCE3xE?S}YXs>Bn| ztd-)4@7XW;uxHd5&&$cj0Qi5U5^H`2?bT#wW7v5uaYjXb{z}hO@bAs9Cp~g+RJb=Q z+*=jy?F#o!g?qQcy;tGhuW%n!xDPAbM-}el3inBc`?SJ+R^dLca9>ooFDu+v74GW_ z_f3WS*15U7XFH>JTdg$pwXmfy*Si+Nye>>5%sKq8Y;%47N0@8Sjx$faw*D=ipSAoY z%+G}X6z+}q{t)JMXpHRfI`q5vevto7nAhN6g$IEDMR+6N5wgd*{Ihu8XZ$40Yxq#f z^ZBqJ#q&Po2Vp)R_PsFwmhqi%bJv6?n=XV?-fNj;tRt;gYF@8Rg=<^k+Euvr&Xv~K z6Lg=$96J>3>}=ucuyr)&8(e<=!X$%>=P0=T2WQVwG-tE){qmt_e`9M-FaADN@@Icz zYEG~8H=^)oe+Qx-`)b;irsmj#(f>YwvF}Ykl(;*n9K92P6&o-W3gXYy-a)R&P(y?hqNPvx?6^gX`p9DT>K-cQH=4CANv z{wlDeik!uMQq{>b?c@ExePKVT-Vb^ex#xNJjq)3N-^uAV-Y#OgzZjR^!Iwc)YFD3t z@y~8AduDRX&!g9SqhmaeUe~R~%|^!Ctxx-B_azE{KDQWCJ-~GB>>Kpx8oZ=ozMgB$ zqoAkNa>I{hz^mU4mIke{EoHiNJ@-7Yw^HZ-a^P(i-OHM677zO=)#v4nH(UGJU#az^ zaw`=6u?JGWC#?itjcvukCVLR7B?=;>!`O_l@pIBLFD33}AmRSnDZu+LE6 z)%7cQ*;lCF-I|Oy8(;PnYMiTsSAAN|WV12jS?fAxq`Ynxu-s;H4;GkzSIO#{{eS9f zD}2-YPYY=Mj@@j&XZd75pnjk2AL7)S8W1pzlf8q?L6+OD!QX+#$eNvjV1fA z^j*Y|!lw6SWqSC1Y#q8}*U&-_zk@aM{SDMalPkgJutJZ2Kg;Hc_i$ytvWH7+X?S5{ zy#nLetc_}g{auZ)tY6E=-1k~#G4X8GMmi@8_O#yDw-~Zqum@~4soBo$A2)c4Wg)woB2)*OvAU0MBje^-XZ9xU~5 zi=b2geQ3*oX)Z6avF*Y?+U5Bgi7{#RKjL!?^yt`c16t?k)}Xbfw=%m~&e*@D&p_LO zSD)dw1+DMS#umEb8K66YS8IC*&@B?j_NFVV!*ei@S|d9dZxj?08gEt?-WQf(?B~+;Nb7K>>DyzZeC&5C?D0HI>deNKXG7{- z*b{nmEcO7cKI~rDX)5M<)@$Ux^K3wknUrHyA5wSAM>ZaOKd651RrsBq8|e0V7NFMR z-liwB$$lcOr+o^W-anM-SqFJk-}f!_@Z31fozBty3c39QhA{uV4hR?l&&PoU=4ZlX zb0KZJv(M~&G)rB*AIc;x=tJndL1gW>!2d8mD`e&KT~gp zvC{c-81!gthZcHQ2>&KcS3I^?&1RO%RSP+uH>P=3o!R+fZqLsb%j)rT_@eJ$54V^y z-s_Nz78DDlKCz!l$MXc} z(YiSvwARgW!S1nUH;b43SbG0GDaf8!$nvbRCh>P3G|3>PV_m*Y#@e<GR2% z#+!{fd&~5mcOH0~Mc28Yb==MYJzDUyO?Nit>_yZ0Z~=HVhVxCV6jc2^d5-h>^^!_~%{<#aFbYEG{Muh#ChpjFp3rYpy5V`v-0`*(^Tx66(CY6}+nJ7SQ7)uQ!{nEWYaO zyxR&nKOZmC>3xM6cBJJq!?NB$pMP0zpkrBYpkv+}n9XVS25SA?Uc{UI+p*iodI23{ zFQC@m9fb|{0ycVHh-cYpecxH=i4fuM|VR{KeF3(SE2WwfVDjun%unwc3*+r zUtkXu*n2oqXjl!sBg{7V};zSZHvUjd$`IyUdTOBU{4mPfjv`T z&lcEo1@?S_y$~>6*B-W%IDI(0bp03$#8ny$O09@#hWDx^}$|TJHg`f!62w z-CU`|K_i{9e?ZrmKZ+RGKhQ|$1K~bT zPWOfX9^Xig|Nc&_?be0=r}u!r;ft>Ee}UGz{Il@Y&#=jA$p0NW!?rYkvN~BGnA-2K z#6P}gyj-ueHoAu#=)JF_>2ZJ7rdqrH6EL;kDPVd(?;J3_mv#vlP5D^y-(z(=6yv|Y zHt`R3O#(gYPgm2M<%9n|tNu+Ie9#)4EMSA(AU-D#7$NuF|Gm^}NiqI==@9r-;4w%Z#5-`ortO3(;pUto=X8v2I^3ERgXdKxY5+0+U zA>lFUSj`c1>Uhjq_{M)T(|Isg(4)Sc*ETr$dcK%pS^WIBH_iXtg$+MD!DCeWzo3?N zE!KPaJVB?fQS%n@^SAT{Q>|U|1zsKhZoqUL=QrNW*Nvf9eI@1NxJ!}8T?;Jx9X8kL zznPZhpTC)=_ofAc-=pZ9U7t?R|1)unHU92Fj+}0P!2;u6dDW?M@i(UQyYfOor(z2i zSdRi*q`(#}u$~3BSb_B_u-*mMr@$63uq6ttZ-FgYU`rL)(gn6mfh}8L%N5x21-61= z)3LyC_50_e?Y}L+w3eGe>pO$hLAMJ3YM^zlH-Xl71RZ8*|Bm4A>D#{}=m&mXgH{Et zbMO!F>vPTTpw|}rzk${@d=-RCmh6POXGd$SP`LRAQt>^WE z9)0H6AYdBDhQS8wz{i05f!c;MMuP2{4_H8wXy>yWYiv-c60y^P|51H(;vg zknp^s7|%FoooBVREHL$bv%q7>(+J`zmza?o{ZR%-)k_zWIBcEXKGuVw~CNLbvSl7`YchbGa*c zHJ7`9ZgwGjjtAW;^vi?eL)wO4j zLYD8Tbj{xzygF8UnQZ1G-(#uWeF|RgUDq{tU*pa4;_tvREZgVq#)BD8%g|CO0 zKiN1>38h+#2brGChq$Ny5b&z(!G#Y~wK<>PPY;EET0e)Hp3L9A$dU3M7Hke-3hg>P zVCw4;0n>SMWWcD->&ee7&gyqvduQx^H@hSFJQ_OHx1$Oh+;h`tz2zPo_4OFj(^+h} zuiRs!bNG1hwhG^IpmjYs)^v3h9@oV^Ia+_G1lf~8uO)IPfgUXMiJ&$96F@hLk!LQb zty7IR%en9W$?BEwly%HchaRn+(+a)5=O?r2=Qd~9dZ7jwrRR|Qtn@qP8AT4bA4v0c zrtxNRa(|GHR+6W3gyXS|vJd3kYr<~R5U69lH za%Z6@JMY;2c`(dn{Z4Urp)=brh(QMexFf{XCmvIdJ1}UeR&eJ#`Z+8qd9xpuqaZi=T{gT+~v!(_Ai zo*Ow+-nRExlO?vUi9emI_ z`_8Z|hujmTyg!0h$L9yoIzHckpIlvD6y*%djraj89)Xxd%#p{7>KZK6V7HW8DF? z*3jRD?m00=>POe0y9?;S;!|hPS{t1}vj#mT9kWRc%W}niQLBsI$-p~A=t&D%?vv6n zngYCPcXE^UI-*+7<(sgpcBU+N_X_npmGOG4s%L7$vRJtfO8uP{yc+v7COfIb@BVTh zl*-Nm-WK7T8MKc5OrUk_X9PV&_-6pExtbpIIKig_t!s+fo(-5jU(IT~*}USuDXrN# z!K;4G0b2c>-E?_w>4oQjXRXtx%ntWU=^8Q*^yoSvItW7+;I$9(@)<_q^`^p;N*Zb1|(7OI^ExI=d7$INN z_iuYxnE2eKSllzF@oX4$>R4?Qu${!0jRQt#Tf<|28UO6}`b`SC?B8Qu=cYkU*U0|{ zjP_h^XmUSuY_me{$U-jud&OuPv}=Uv^n0Y{Y$Py!_Sn3zAA9Y#cLVqwWpbV5ZypoR z)EOmz-vYd2g&tkl#F3W23&XA5GVRZ3{hX+jE!e(e-Y7$mx2uU14)1^yu2RL$Imq){X&F-?OunJs;V5 z$&O`bB0H9yiR{?D7Qa8ss4u>6FRM-N+tVD6LrnTzawm&1ljXj?Av9#yc<_#M0-rmZ zZ03{iH}&zNn8(6BfvRVhLeH)R#@;%eJ8I*zB46GUm-)?}IMu23uv_p&?jhex7m{3fSd&>V_EfIebsp(6QSCw0@6V0sP}#2|h_{ETmkw@9WLRhdp_N$ZXeM z7F)*49z6A7Z{yACi9L5ZZ}$bSK1b|hvRR(?EA(<-Wj}Wop9g?f?d%`)4iVl14a;I% z75-|z4=Q*^BMfm6LCO!8L*tldi*)yoV9};Y~i0_91)A~HLu=%2`-`TjZzfkM! z2bx6@f^k8#bd#%c{v93IFUQrbb0OQ^W<>> zqnO*tes{*PDUlb>pDZ@^G^#%*z=qC)(GdSj2{$>1GAH|#peWV4*H=TYy2 zr-4_W%TEP8$Q9x96w{UY!5&D>;Tho7XYA8WHmhOwK#mhTX9ixat+RotPiF;Qs`2q) z52VVTQ}D7UQs>6G#_RQ=db1u$w|TnFn@sL**l4uAX1$fJGwZE%jJ=gwSLefT&HH&4 zM>b{`V$Aeje<65>&_%m0FxhO(*k7r3Clq$rZ>ekD#lg-+CfmggX}1U-bt?79JTBR_cl`#5LB=PkzT`PF%SYrrVx@$~3TD*w)s^`p8R`wbh( z-B#Fpy%>|MH`VR2H+3tTwCi@WkmYX9`NcKe>dnM zqU$cx<$2fIxYw}E4*O8`x$c3W_kPgh#Lj)D%k2>I+<2dAhOtkTQ)%(s<;MMOTPZqL#ucLOK0q-~lWY^Os+tFyn zHl2}woH)PL-)DoIu1(JwmW|(SSg!|}Eq**-@GjS*{d)OA!TS{QtM|1RgMY1J^QA)1 z1nAK@`*OkibQsfDf=!zA@#i@c%%$f<$K*B0wGg)J)k05pcE8);`TROwq#R=}Ffm_G zbPV2tJ*}@dLF*iQ1GLWN*Ui7IR(WQH=IPWAm`@aprEXE+&8`HWpA~wRGCiJKiuri4=UBfheE~h=MDORO zH}hc`=+(LWW#Pjgi&5#sc}2>7cocecO#WBc{0dmxvQGuLzBb-0wxR0{`Q3W*2!#SyzDzy-p++T>)E`^>YBafI)A!C&uVno zt}a2Z@=jt{rkB0sTIZ92S93dQ&`WckTmA-yj)`LIC2xc+UC4Q_c$Q1{ijSjOyCx65 zX&h4oOkWCHVx~^zrQLLe!Ee z7-4I37qR6S_eSXYw-WRW5WW>b>v*gX>@IJ1vw6$C5<^AT%Ep`dIMC`l(-VImW|hJw z_d00(u3FgSZ@lOlP0I7ay$i%Lf4IM34Z?QyGaK2M$NdGXfp?G#;O1_`5E6Rd$`i4txH!W(FIt z=TOITNWf^)>yp3M(x_&${Q)kQ?GJElXkl;q&;zCK35FH%a_@jXXAUoH>|paS8z1f~ z(A=z7*o^xM)(v;R|e?dtqcU`_USejEZ?*Y|@# z>+`@tpml9J5VWqX2Y_xBefxvf=eqqs>sq!iXkD-N0XAi?(nt!u@+ z;MeDfJ%Q`<%O0S$hkbX@`tKOKf!6uEE9jwOZx_&d&l(R}pUZXzt?SV^(Cdr*PN21C ze-`+wzxA>s@J8$I_JHYl9cK2jHSr^?t$Keuyzp(#V1rTky8kKkG>R$v z{TP??-;c@ekvzXv$Nfn7rsH@-;p=hXKA>yzHj{MdFoN#dbyT8ftYJq7Y>W%x^O%5X z%^w@E9YpT9fDv+gS0_1kEcX1XobT=SX8;}3;|qJATAo}^b!L6z?u+-h>$o{KXyNwj z&55wD^X-J-kFFgj1x)?vhG(W&V{S9{2qj+qoe?np+5Nt^JE@8PUZHY6 zXR`5UpRKMxXTgTfjWY|I{(CMShw3@Ikc&NxI^O4)TsAhbH}+iP&Ek3xap}A|AH15A z^9ot^%IZDsLhz1bdhNQPkY%r|*3U)Y)%YfutmjN~eR056bA$N2#IP(c?2Fa4_tKz; zYP?Q(HmKH#V%d40E|>MII+pdTI@THaXk^d(RXrE#3;R_Y$@$(g{Z8xlE<}8~MqPni z=y+TXTF3ITBCiL*2c4T&7QF0T)$zKj;AO9>@?LGcSsu<0c~IZ4F}W-c>`T>pzaG3= zhu48tUDp;qJY+TFwWl?EBjhx{H<%vh)iGsXs`B1c*xWY6t$H3>tbJ{?lC@C%2=UapCw-ml}@1e@xR`~8cpzcGHJBQCZA*au$cYxkX@Y_wV%M$Y1W{;-U&E4SD zd4HG5dL3!5?+KXd&3Yl-X8esPwRb;s>a+KKpvOxr_nOVD7GfXe1BHKC->m!Ry^fg; z_B!euc(Cw?y^p%)KUDbRy^wBGZ9D=woo5dhHrX4gHXki`*$b(4{+RJ*bDzDBs`ts@ z(-WYz1|K(FUN^cf?2hk4x(=@?Up(m+{&m)S(emz}dpAI@Tm(K)D<(@Un z$3*STj6ApS4|Y9Q@Uq`dbMt(#NfjPr?7!3d&I^T{pAVY(dJ%lpdHZ6aXJik0lRwA4 z1UVhYmrYM*lV^!4?<p)QEecM2bPPO8oY&C$C- zPUr1=g?;uiYMs6x^yoW}4-Ct4zVX=IFD4QH_s-_W*jGD8?Q@9na4TJ$??W+t>Vk5#;na;PS9O>a)?8K~8h^zkumF z_f^2wl)QZ%FpbyGO7$2N%X;lRhq2dA<=7XewX4{qz1m~y=Qo9)-xkePxhQO!oCZ>XuW)ISXNu?DO27bf*!5y9|J~jJonyzmc_*WvqpaRFWv?HRQUZI z=DV)5KNq}fz=rbv5_sEg1#}VDZ@_dd{ndE0xU%23^L+5^Q^xE$*P8kRI`!S`?}ZKT zfpb0T+h35=xc)5k>|WH)slEU4HH(owXBr17kC8oQL-_~0{x0nE3=pjWT_-v$)gB|| zHuewye%FtnT^)m*>imyko;QlPPWGuOZ>NHneQNse2%UpXT|>JRdf3yZW7pMqvwIoO zfYmuUDR}jMI!W+BV`I;o=2KZjo5$B7@i9QQXSMc)R9@MAljW?^M zKDI7pv9Zre^~_M{@jk0e5Bsb%&KXTl7pB}UABWgSH7j^^?#%*P*VUOptL~XhcNQ!A ztQw86_fchX?4xQV$37@x*?08V+cXD!(^{P!w2tX)h41WrqMa}S#(%vtE+*{hB8 zu%AhDJ(uap#)&;jO;lyqeBjkunHRLiIuB^|acB@ZL0~%Px*Kn1mpx1xe-H4g-GxDG3=5equNlP_37Fms7Y!KQ_4>cDs8K&h)$OU= z=63(eF!nYzs$ZUa+Nd7b*Q7D^G`~I1lytxMfga8KVnvKupOfqHKBvsaP^?vLVbUI`$d$feo#R-i6J-LY~{c=FkmXiw8NymME}8+6HvuJ2sW;TgWY0 zU`rL)(gn6mfh}8L%N5x21-3$gtyo|y71+uJwn~AmT44POtf|0OE3nlKoB5yeY4?tM z7%*LXhl17`9Rga{o57%)O%Ff1Lurf9>jG~Tx*KqP&$$k8edjSBaDDDy8~8wxUkkLZ z1%p8A8oVZGUEc?S*8ZV2K4iZWECUjYnAr~Wos4te>CRx3O(x=SnTl`O$Y4SAjpky0G}Hg=Dv*)Y-3+o*u4+~|O* zy|_<#2l0E0Ah(lXTN>u^ZY|i>z{U%`RnVi@HUZO^#sq8+(K$9?gnSJ7{zvOu`}DRA za{A7AJHtFCov%9p8%DKuZ6EaL`<)#FMiKXIzVLoS*SMX6oUS+H0;cnD=YXj%;{&E^ z<}Lx#e!^W1^XG|v9ymVt0H$l*?w~b~y9K=($DRRGy?X^rbF;T$J~nD^UtsF{K0%MJ ztNR74N#fc+V1#_!4lvB^Yy1ZRQyT{cJsSVP0aF{j7bZ3F`;hu}DC9KeLrjnRqvLuw zuwhhX*I~xn(P+)<5kU`GJZAnrx!!}040?3zjtY7dI|kT5v(1mA1FznjjsvDPj}5$f z4?h8z-t&$(UXN4n=O+e?et4XpElO(Q^Qb--oK(oM2T|onxz2dzt^ddrQS7vUQN(@YIksA7ioM;v{d3dlrpMO-<>k4xI(BD*S8MbPll26|^wczx{D&aQn)P0Wd6ALILp)=RwaYfjDy{-~aBQ3Eu>i2|zsV^4=OzZ4o!#w{q@GFluI9VW z^=p6)w7>J?YS3E4SAo{?x)Ss-;lBd3-oGvf-DvFDkE`)sYrI*ljzXR2nteTZHOJSP ztk<9Fy}>Yjzv6l18NRyq-V9z{qizDN`MA+^Wn*#>#zc81=e** z$LH3dQ*(V=z|>y9pi?pL;q{o*-W|}XwRF4L$i|30y*dZ)0h6H)*!r`% z=V~0Cx9grDr}@7(V7k`b7ciYm_Z#LhYWxoY)B1l9wBExW2zu$Jk2!mTIj&h9viDc( z@KNZ|I(!7QKF>aEc6`j#4*PO7w~qy0&BGJGG}n(CujjgnitT#Ru*`p+Ijdv;40v@8 zJ`H-9E5hefrpt3o$n(UWU%mG~2VNc5XHC|9)Oel`SXstDf z$#pRsD#w0dwec!=TZQix(E8oxWzbq%>sr0|7^|+=!K;3}2AZ{%`N4i-&BNQ^Rb6id zU2lTc`gy~2yB}Kb?*h~3fOm}7$EKge`(D5Z`IvuA4@}44JLB~+ z*Jp+w0;cQwj{(y-diEwYF~`j$?fNOmX>NWt%wtr0zXH=|?O%c(s`os`-;vQ8`>l{; z-?sX#a=#aH-k0ru>m2?Qa!pig*B_?GdHHAeo9Ef;diR&{W@|qCvbE+rEZx2a{w`$M zr`_TX;qyO`)!Oa|TJu24ee|AcuNTGqeAz5E_E4+8oykfb&%zs zeJpsUtgbuLgI9f?&SX6oIuBes`QL zV4Bm}1Ez7!5iotun=@dVpSc32x%@2L7xbPzcaT#*=P}IdLUqO-@%g|zTyit7$@=)H zEPKW^uKA7EbF4MrE$Gp4?hZ_O7YMv+ljpu_YzqdP${WvtRqQ*URR*R^Qbz&p?b#OLzBbRAi)@Q3HjHmWV2 z2TROz)k4y)6`@ndd4=E)**vdT!A7eK;&UbA&1!EK*worv6}%evDxkIYRyJMPoQh|* z>RQn+$mtw!GR))F^=NfqTEnXaJ>w+K=74EEw**XMYz-JydM^71Ovhk=VeWUc_%aZf z-XqosdbD2F44BTRLBWRlwiYnG53e2cXpOEDFs-e1gAI*i2rykU1_wRr&&#+cX>Dtc zhXxyp4G%Vk1zxR(^?~Vkne_s%`V-H--N1Ogc6EF zwSG1YdenEGk*oFl--4HC<_@Q@U7H1){Il2dhyvqY2pvDo^X7&9k-(ZLYS*ZOm;V;6 zH9NZC{Sxcga94`YEsWR4kC4aZdl7stYJF`5Ih_kz7B+cSuijg>E^PYVfXtul{8{%k z?gh|im~DbR#l{3ob3V3+=|DV}=<~$3L65FQ+ZB3vHm=Ua?Ty!;U=x|nAXO@#+%vY-UqdFIC%Abbr@)!`-hq?A770(+t=XvUD?*c zEKl6?pm{nX*yu+ib{!coef~VEz>YS|?dv_|SYTRT#{@kz?Q!_;2X}IZ1mnI2oma<0 zkLK>U;DhG*L|{7SPAGifzJ^Bgk^305cGnB{AdT~+U_=$t#nc(eNA zJ_fDZ)4{9$od$Y1mD_cy>FVNtz^B{gJ_prxHh4APXMxsyp9xx@{mv+KbDx9O+4GpY{F7?}Ee5onzY6F}?p`h})zk|aMoGgf1`61?j3 z6`*yEy&SaOQ!WFowk`!dT+yfysq$vdo*<1Z!lgTJGIF@7s`7xcy%mq09DvD|08?t|v%L0~#}9{{aB-fyy5{(O%_hVk51=GJYh zf80Nzb@MQ6Xx%(yHnZ54g0D0(!bZE1@>5hozxc4#D8B>osSoCPZZdb1@=^d zJzZd@7k-{mV9yjf&n)DgE##gnu;&Zxg#tUL@a4rq?xg~IxxijAY$mIqe*acdhYlYB zYZ7|rS=z5N9|G4o_W|f;;olMbEkeH!T<=xe0oS^I54hgD#sVKG^6vuIdAkj8y{~Ky zTGz~Xz^`-cZP5Cx@D^yD_iuvMwdoDey5_wOTJL$A!`^W5?=|4Me!dD?@1>-Cy^80c zDz^Soq$bWq#dr>CBdja#vyHGx4b$~y8?xj21b%M!&;M=v@5hPX&ifbEv1r4wF}SZb z!UhLS>*~|O-phr(r%h)T)929H2Gb>6pA|L^v2{Ds(|6(aHSk5yqxY3B14g@^8}>mr z%1xJ$8=ZsyD|G%@e}*cp@%)0jeO0{QFUMAS*aSwi$&gM zFu#2aRgV3UjeLn`nyTE8P*1upe^N@>Q&}_UfEONu1JY5%l zwwUz!)ot+HP93XX!K?SuUy6LhzB`2Gqp=rhXS0n-@92c3##zcK8-#4`(3 z&VOS#i!b~AUe_7V4^*9-_8>JecE$Ysz|7vrLXPMBHPY$l`ei!VzoByayx3ux_82LT zuV-N|>!WbF*hkSyrFL~J?8S5TDB?Q3k0OhYeH4r@TMv0Yp4$IUu-_IU#c_2C82{$> zz4sxr&)?P6{Hoka#hhcWK_j_%o}KE##J6rS)-UyJ?B|tz8e>QF zZjA9aQ8o5aB<<=__}R6wKG?c)0&(+VD!fQo+n_v2-3 zu)r27uoH{e7cS&_6xbpKwrGL%EU?82tXF~cF0eiYws?UpQDA)wY{>#!s=$^muw@LJ z)oNMi*qXp}%^L_>*UL3P>odv#(9NcYAN@hM2;B<0Rp=Jb`hL9`wBAp8z^OHae|6yc z9%nVs`hL9$w644TKMMZOd(HBo z*AxGj1FgM*%YxST@3XcI!o+oiqQ3qN0b>%T-}d+|5P z^x1Q*!cTvvuYPV#Gq~nY-uF9pEA9(B{lnJw-?f8{tpr=gFu#ZB+CLcB9>TkB&_l>= z4o}{JI`&7%>2{)XNTIWT_#UM8{CBoIH={-8(4f=(Aj`0TsZQUc?m88FFuAw8-r+&# z2nyP@UcmTgmy7QTHHP&IIe+KpeyY6B*skR*gV)!Y-GSx&n@91EzV}Hegig<1``6k)6ff zb|&ZXsvO_Xsr~H>8}WUe>WS~6G#5J*I&Tkax5m3;As62Zc^>G~PC;%A|7h2^fT_Km z14a?oyIxpFb&Vfi$i?1j)wxcnZLO_cf=>5?^1B91$9lJbslDu+Zr7WgQ|=hg5RdnB z$Nb$uR$H5(CN)pHn{OE}dp~t7_9)`;^N%~yExX(we)m;h^?ti&(5d&ry#hupuCo`u z&rxUQ_dbiX_j~W4M`PM2U^++k4Vc=G-(OYdekSL>jijJm2LMwW`xiD244B$D$S{BZ zq1d6ows(X0JOs3^p9crMLxlIRfVBvAc)$qxcq|^qSTR0->G&&lM4@xR;u-45LhdNT zJpXF%VAPGStw#qviX9U$)$?Hy@1nu)wq=UKxQ-1v`~CChw*MYyn2#y{==nVt_F9E^ zbI>|&#|OQGh4%!*vV5=?qeXa640;H;4<{9RcpjX77dqK^+rHkLe%Hg_uhpEK244Lh zb}DF%?G)41QT%jW*?Dq~#q;EJY)>zAW<5Wy^C(+)+!r0cGlEY2u6w3o+1RoNMRRvH zc=dbYSwXMn?i|CivArKPqI%B_dUR~Pzr^)uT<1e>J<0ufg`Q5~{;4*;H_UA)w&_ys z=j8>#hJFvfFkq@@f?=7BC6EuT^@|GLZs65CUtI9=yO)mjB?a%O;KQZH>-D7mT^2BX zS8{n_gTIHYdaf|utoC>&n#Oq*cs18o7Is#JpuP{eI`FCw*8Bwn4qnamZJ>1wZw0OE=`Em#iQbz*^ZwjL{^l{T*GT=mJLtX(wAR3#pmoi; z!*pkHuvba_egM3h+xtQ5ec?XPdSAO2wAR!;pczLtj_h^Pv3Ll)trG8pCY!ArTO!}O z9zPuHXbg`5(->VTWh1HHlA;2VNcXCxPjhK4HAse43&y@cdbseOBu4Q>G`Y z1-@s|T6h+`8r#;;r9M3aT=Vv{>Fz2qdYv5vS&jWA@HPwIi=frd7eEgZ{CUt?ch7-l zEo8R%`^_2$-w~;;*Fg6d-LHbyI9~y+HUBbb+RAE&ePp|cuGft>8>cr4%+FlQY&--T z>ff7%KeOUqrFnd-u=%!O*?i*fJL`OU7rZ+5?-V{90@-HC$$JIwlrW-p-Z$PXe;eC( z;_R8TZ>X=2pht7|A?WqRrw@Wpt%7|5OzZk%<8_~Oo_=armP4Kur@WtoSJ&mwOg0+1pOXJP*#J)eB8()K0=h;`FHTM4%b|1lWt;&95yjh;U4H(nyYeTQz zq}u-3$w=0d=<^JpdUe)^xcy*k92CcR9 z6KK81{RmqB4d@5ZjqYhYYfg3l2HqBT2%o>2Y?hmc;V1QGHEc;nf1e;aS+KhLgfR1;Ybs*ewUC!BI!_+gi0>`&FP>$Gfp-~R(z z?=c-gujfkm3A&L_e&$^#ve@NzdFGwk?Go&EHoe(c#WV3H1+P9MO#)iSs;lYBa?0Nw zr~lbF@;tq^l$-C96?z^iYU?oMUt^sTdYVipKc)b!{!eapvb-$>U9BQJHF$NNPF2W0 z4q2_sX~C;`n#N?aJRMrBE7L(v_DFfT=$-8gG_|9zELUVHWVJu9;2N zV;Lejn>AoMHnSP#{$8%eZQHwL;bAez}^>h!|D3Mz*V5)PWU|-ja9>4~Qo`sFqYle`= z#d8nU_eH_0_r*m_*2hclS3Lu!dOx#qcAILWH{`UQdV$tDT`cI;*!l!a>u~XasXu-O zVm7wysaN|;6nYj#z8cl&d2Np7bBMk98po2bp>uZ^=-r&E?doejc>efjA2asytL)O? z9pHlaT*_p#Sb1im=H6>8<6R2hVKgtxm|m|1^@%wX*6VRj;Ik#@*-fz4fD!UI-a!s@j1}u|a$WgHTsr=10n^+K0+%oO_FwnmZRA+%V}0-r5We+5tB=D$uP^v8(2d6E8u&;o z^Z);yUl{bLPa6hI@5LJhOy}3ehPi#!xhXK6Pn!f@dhR*;f9Dkjo637!SSK`hn-y}; z7rA0@hdvLED17CflqNc4*XG9SwaY(yT+d?-)v+E0Uah~8g}?B z(Lt|%N8i#gA2aor{V#fN+sb$|A9$Xh&W~-t+hRFs(-~wnXWIhPS{VykeH;^H)yI?I zqsnd%UX6XbLiXe^9y@|p*Onbj*5|m|y&kfS<{I}D>A3D>da@jUh6HQg#uYYwKT?)| zp6{nMFdll;hn<5DYI9d$D!YsEW_H+5q&2)dcy)c)t&ol9{Ot)|9lJeD*7Kw_vUlL! z%XqV~`QFAR!}clUrob~-qqTLeB^zUl_WB3Q^3@0PLg&zau&>WD`xgHAJ~GdTj@bc_ z)A_Z(>2Y4o>p{Whfx)KM$-#zYvAl{JQy&fmug;f43R(7=DevLn)p>fD$-4i{sIQx= zB8U3#EJqlxkFV-EGGMCbbok$Be6MR8r2Ks!_sQuP9t#_qhhspuQng)2gVyJ|qs*r) zHtw_2aXij=JvOTJJorAlY;1TgqUt#TdUT#1AM`5kiNQa87CQ-;*2ku^(WCqw`wf=E z%=cStPG&jigSn;8@F$x;ooUK0&jI%Z>R6lxULA{5L2K=wV!ATB^jB-=bmPtP%X1&q z?wR0KpU*Ja$=qRly3afZQhmG-y!za90cc%+&j+o3od;U~z2{ue`i|oq(7NB_Y|uJp zXMyIJO_KP={gUe2W#DZRzDq%CJ}&{?EcnHsTZO&|wE8{)G;Mbk9?uK+U#jk_!K-z7 z6===Nl|lCvp!L1_<)9nY=7_ez%inEqFQ<;}b%haHwRRo?t-d@6TAzy_0IlPyaXua}a=M?JwRH2J!b!zVRfiCsyS;*>h(=(IZkyw zXS~_m@N+-AFhO>Cp4hvswe}M9sNNSrtB)^$Uf-4A^Lfx5pUg+@UDmOD*?8T@HAK%V z0VCu-P9M$?)7pQvkUJgsQLP=->1TOnW6rZYH3x6NhTbn;2dzH5W<62KP`B9CZY0vX1rOB zmI0&Y=nL?+(y(2h7k0h`rgr{kyxE#O5{%0G19){DzXz@N_U}MzZoUPr-Ud_W%KGbtOqR~JU)L0ua4PIpjFpl7?(zC%kEv<&!a7j*XL2&*HyFi zOYm2*UkmKF0{gwd{wT0N3+%4~d$29AiQh>yhQABB4$J+M!z?cy3+z7y)~Ucc7g(19 z>w$TtHKp-&E#xLCut^JSvI3jDz@{*4TB|Lc-!lW#`@~G3bso+LTIb;mpmjb>4_ep0 z=|B%KJN&q5w)T7YrvdNdUQ6g{z_0UuYS230rUI?sxh{je&ig5WZy@?E z0j^_4%I{s-S+$N$67J_CC}`I#u%S9Nzq1Bxdl$mzYysO&u-OAPTCh0+rgb-Gz*Og4 z0i(*TOz-4a=QM^f_IG~FUC71r$+i`_d4k;Lg3TK+LhegEBTKP(M%5_MIbWf3_qLnO z#I;D*@c9e5UJKK}L^<`v&*AlTO7BPA3Z3!%T-A9}xF=|>FHq?Gy7(^H-Q@Jy#dFOw zYsb+6yA~{L>`>%_dvlnl%tkk0daqw7*wC?9IAGN3zGmm!I>z&DHI5#pN1t(=m-}mU z&07?_y5=onvL1{2z8ElFzj_8DtuGc)eaUw$Wk!segTfocg+Wz|^-T4D&eD z9{1X4Tzv~(?yu3@E*Wg9-YKjOGn*&F;6O6lwN#Jl}&RJU;GyX|zt>4IApK*7oYKr(?HTVc++` zWWI49jM^aOzVWrqR=%RMhThOGrJ_pE5G^e=e1=R}`@1{iM^cedxm zKJTT@J<6>JsWK!M)wwtV zdIqp??Ai>p=KsHio%0Y+qxv5YJNk^bx#`J#VsEheG&=Y+3bfwqMw+gyM%f#z_vtN+ zH_Ma%etI@mm%x7=tF53%^=?_{ec1A!`NuQE^;vuCLQi&vxX&T>6t|_{#%&wZli7R} zyz2Xy!e;z!bEE}FUzVU2uodeqjJz5vr1x$VG*o##D{r3{ADT=Ai z?F*f;zgX*G2b0TcXJh24g(mIV3A~#B9fQqgXT;~YfD!Wi-&Dl0bC4S+I>!f0b?#DN z-j|%so$MTS_w(uCr`r1|)R1D&+I=vqlh}W{Yw%mgX19Q?F1gq}U}}FBtD($(E9Q`{ zL3vNP8`H<#L)ndS6_kQ*JtJ%c@6L-sPv{nqt-A7H9y@1RHXwr{{RHyx6H@AH`0 zm#Xti<>K#BYYz7-?8Wogb^h#ca@pKtPo~cA1Hr2{4+u7!=&W4_1&n|87}+1$=w8R3 zMqP^!4tjJgIK(g?E49J?Mve2(phtPT7vqtg_2zN#TsFp$#!wTQ| zJH)z<9Z~SImrrZ-NaM|XV(*^T=+WTS935q{*?jVzJkPc2ITmvIKK~fgN9 z=s4re{AG`w*53)>)%rW$WHW!W-zN5P<5~CW|A|45t_LRtOvm)(fN7ku|4#ir#pJU5 zus==hpIX>pzgVNSg8f{LYBQdfe>!aFn4DJF%+3t-ybdhp1N*Bq=V!u(*7zC0W)t1F z>#TtB&u-JtDa>*?Hax4H9rS3P&k2~?V1HGk_3ABPn!9s@9*y}t!?Ndg_EKrBUkF~! z!3CgI*ZHQa6J4{*bHd&y&Fw|t)p48) z!an<$)YnUmH;a>fO8VSXoX&$dNxvcis`=jyR0A9`2^`Q0NFs}lf!$tUcNEy21$I|~-Cbb!6xh85c3*+rUtkXu z*n2oqXqU@fjwSePZZdb1@=^dJ#E;mR(Co_ubjR8+I$B|(pmpte9<;s}JstA8zB~t9=l5yA^&a#r@D0TOQ-E(M^fSO4-G}_Q z5VG}%y&t;nyIcWXiT~`zMIHxNp@WGw*}Pxp6u92 ztiP&rXOeck7j$Ypzi*i4&~=^-nP~r|bR40UIm! zJ}IzI1E%xovw&&tKQ}DP1%D?{dA|sHG~U=lv$gpBWsuuRu>S>2b^aR0SF!AStmpo# zLg&{3)7*a(FtvAK8)IIF(}myt(D=TCPOZytgNG%etIuV!?~aQ z{c&R1oZuPQZ6T5c*Dq#2HUH1xo8I^S3K%_g``K^PIQDnA$LPKF@1RqkJ31`iI2Z9Z*439?`?i%yZ3Q=fK-G+*FOLD=;0mE`>ka7xkX~=DClB_seCz zvR_X9nWV7at(c4cyFr*(q$&6ROYdbG{?Q~400I%MoCog1U z@7+}3)j2Vx$@-eBKE~gr*XPiw3%Tq!>pd3lyUXI_S=G7*O$!@3pQbUJS?>AUOZuEY zUE$wkR_~b}_UUPUrZ+vAzdR#bpY3M^ua5r=g)DpgbWCRkug;H|Og0-A@6~hvbzPYi zaysv4F+I+!MqdGh|hEdM8laa|C4w6?kzK6Hkx*3&|T58fY`#m)Xe9jk>4J?s~3H0BoqYx#F1 zI&|n!=wW}L-p5G!n6e*`Se7sL|7m?L3L82v7Ab6IeSL0oY~hdh?qy@k^Qjxfk@e=e zJ?}rv#)Ex(tyF7Q&myL*kI!vneSD6^UOk#}jAu7D%0JIr);ZD(zG>|&X1-=|@%&}2 z?cN11&sSF7J|PDS(M7u!57-6{;B$$9sm{Iu)BE(20i#OKPj>dQV?0k;Z7fyj;d#o; zbCxUiHfmm%Ha%S^Vwc;Dy^hO+x5b(8xg6*rLN5zieO$(LXZ7v9kQw$ieAL=+6&ouQ zI{kd(EN=Ep>a)v=rYGa&dB?gItW@yw{A10*$^|cbER}baf|os)`V6wF@n-q*Gl(-R z_G2=YKEC}-XLkSNS-)`i-)nX7HvLPXhCpjwHx>RZ*Mow2J+kjpbJA>jvNd|c;8QDj z^*+^NvROP6AdCO~d#w><2Y}Yv?QgoWoUT^HGZ>i8uYsl~tFy;2Py7Ar5#9+GHqY!s z#yp1?AuxSr7*yyPS@_O#UiH27TBaw<=lduQjdNY_HcKq)fY#^uwM|!+huEV!*m$$t zv+uM~9@r16H9MrRc^(*B#NVL>FZ(~WMu!Fe6dP_>mNWKqs(%}RS97>NXwCn6rYoBZ z?Dte%8-lmL#N+Q!yU5?RwExe3P?g;@_^}CSeYV{gw2tLQpmlw2H*e->kCAe_yM_In z8vB2PPL*SysoGP_`%1l5bSyS2bo%dnX7$J((*ZPM*N8&TSAEHpzn9}#w2f+&XU}R* zHxD*+K93BT*2t)U(W2*jl4bs}pVdI@LEH*@bUtheTJJwwfY$HfqYJw~TiiY->eF=L zU98r?)`iaNZT)kdT65bJa_n_%G(K5xqvx)Fkkh&vQ`qax zUd`c>g^eMwp*h{Qu)%XGnVW2VVb7wz@7^AI)c5TQy(dAoQO?=Z*a*vhyVb`d`|Vc8 z{_6RU&TKA!fn4c)*`e^8J&WpFKazIsXmV~tv2nn(?_{T7qsf)vbLW5&@;H`;uWDbh z$?(0U5$0zjX02ei53`xcvR6^Ro9qQ%U2pa* z{C&yhQ>K^wiOo!}UHd?f-q-dvy;*Ire^Gh&4ZQl?y+1I0_Si4*st*SO)7TCOymZy) z1N#}ZJ`Xb9EY{fHcnEk~J#c&;Y_i$7&RNu%pBI-s&$0JW$M|sAP=5~tt@rUmgAeP9 ze@6t2kmtesAKiDgd1N8Sv*kFqGT-CbamPT9uJ=cS)^R$ju)}^w+R4U+{fw&-w(B_P zX>kId#|FJ>^Z0=2{5m0ERO$J;$>R2WDYkimvG=jjTEt$*Mi|dnQ>-6JyG|_p`&Bt%E_>yIX$Vc>1U+*oY%2A6>_?6o>J)H*=O`U^N;GHZ5=ZyF~Ctx}k&Na+q)z~forgQv! z(E2QJUeK#Gd||+J3?~E(x3qt+iwf-G0=uNZ-fc_D#P2a`@6tleAYN5}2+tPk`3-;&IUW z9Pk)uz1KeqT5I|d&^kXK2CdHx4}sRXe?Q(OtnqKvcj)jS@PR_#3tZ>Z1Hg6tx*NF8 z!j2lEQY~ups|WuD8YR<$6DMy|+MaM^}Q+SIqZ}muFt;_`e$B z)ZD!m^yoV91~8q6uLoYb(_e$Q=R9) z??&UXR)IZT%*FY`9M*VtS^ZrSgl5#(@ z9#F>yE!-S6ZC7YM*oz!{ozk-+jlA7wS^J@%V`sW0CK zO!L;w?0NiJ&)-8%?SE%_oLBXXf*xIqekg2aXXAPNv0v{;)053V_UGxk@iTaJuKZN^ z5c~5qmP10lD7I>`j%2-i9+&s-Wpm|$Fa~OmJ)0WiFGWo30cqSfoYBYQRrdc9ewS>54+q3`|O(hkNjMr{rvqK zdi7rV7ijhWPqUlllYMyleA^Ma8|93Bc{-2(W4u|NurIHX9`?#H=PU-tS!a%icNV?ONDjFP*L*lY&?0=OiYZ#o8INjbdf*oaSsY)6>}l!sizD|Jh5Y z>(Z3Lhbcg-uE`6ZvR*s4a}eyPo~a5w-cOgs&3-z@?U+7KOkL>ov+lZB-28Cg*b~=J z{yshQ=(tP=x>@jPLF+h8Q`nugxF>tRo!6k&%}mg#^J_-Xn!gzeo8Bwub5G@Fft==S z=AfrV;+)m6ENASE(;UqSUiE(t(E6QtcF^j_Y^K|Nq*)&i_R^`Yxs2EIq`cf;rt^9p z@alWAxlPvVMW1cv4Vda>U!2yN#xkGDW%DWa!7Tt@eO~ScdK^{THGkm)`{WwUarU`s zKD!rsyzk9(jiLMZS_pDlXA2g3ybsQENhLmJJGBkM#CfUMdg1$nV(g#OJnKAMIQXJ@ z>tUG3q&nH-rsK6p(4)LrkDJGo^|(1U-LmcdR-KCmd(8yx>RI?b0eL2^nbX*mi8F4nMYxb_G{lyFW-n*97 z;}V73Mul%%_WsAWOy{O1m*spUuojBiwIuB8yzguNWU}m!(|TPRy!w2!REVX?jo@<` z!?JnsK59hg)3SvRJfl(Hi7pRb^=Y}_16}Z1^xiqonXcU{1UX%QR}7fye{zbq366}p0IDOm1gW(1$tEP%BDBVEqmQ`-CQ;Jplj8zIcXqY8xJpP(a%WC?lbIN z8|aSXb2Zpkf0{xp{fb!F>!$A?RyW?PR{4AHT2n3H)n}n*ll5Ali(V&7L$=YF^0(o& z-dYQPd9I_*f&PU}@5jshVLx6Yf4Jwnt#rtOYd~Sc&x6dyX#jGo>)D#nqct`Vv_6lm zVRo{bzF0}cR$*qCdw5~N@>lF6lnUOl~>lSi@ z3tzIc7CnZm`?U8{YiDp_Zz_u~8;hCX%K)miYe=DIXux!wh8dRGSPz)yba>FCb+vxL zw5B!)n10{eFklSF^X9)*@4l(MjUlHtHZnb3JRzdz8st&q*bKa?>%XA21~vt)b-sz| z_8RObHbxlcJ`!^OV=wCFL67o!|7fOXx8VQCLeJFTZ8RrF1v!mlbP)&7!_>8Xi=ap2 z_-V=Z_7!^$Ynu9|SoS+Do=5LJ_57>cun@0e>`QHAFZQA8bHSDtUnh#&<-YJ7P0icZ z;MLr06>_7`QQHKJs$K8?;X8`@Gp3L`FvzLSPuq$+4uROzU>N3 z*Unu)Yj42#;5)6l4>zEe8s(XNx{dr}U#{kIw_ro~D6 zoAlP>+8qATch8q%{I_qt7wru_T0?skKJe_+CMvUQpWuVu=l2a5Mch9>(;Y{B7`q3LAb_Y8Ds!i#0z77kYShs`{$_#IsUW z=OINLJS$b_$)SaPKO5EK&^7gN$Zbf6>^iK_ll2384P0jTk?cDJ`+@adaTILmcpV8^ z$Nh-lgWeyH4jA2bA7VeSV&3zct=;VT<@jX2vxm2ZuwBQ&zK+qcpml9KrtptFzgh<> zcYGmtq{(%1C-Aw2{eSl6>Y8&R^yo9g35DLK;E#^cNhaqtLvwDk8FK3T$%T!5f(`wi zeo7&i^#Hql?*YzyV-K*d<);Q48sljJ(-^&vH>;PwLQE>hej^=Yt@G1^J+1jO0@gw` zcAaTh7yko3Jr_5^7mev$@HPqGIiOpGJ{z>or?X7A=TiMWADBKzo@cySPVOn5BQG$y ztbW-?rR&iI@M>NzG+FmgW4RcZ&c%z2*Zotymjpc;*JZ$TExk18(R^JVFtu?-u)(zZ zII$l~=joL}4{nqHURM>^)dhA|H1$IM$-B@5Z71+%Mc1wYM(>CZ6pRqNc zw-$1@71-^DdF^q4d^{JyI-@r3DD>P}V0RVR-34|}f!$kR_Z8Uv1@?epGh3DFn)5s` zU2~oTt?S6Mpmn}YGFSV5kA4Pt%fHp%p+gtoI?p?U)^|QngI}L1p8~D#{+|S`-zhpk zUe}2yfa|>d3*UKkZaxlNpP?TEt!vVwpmj}n1hlRnze1lr+dm9k`=1{It@ptPLF@gr zTiZfEv6i*&X6-{N$2u2laMt5I+H&5p?S|Q2r(^rI6?u-0>D%6y5vGqHFCeDqr|TTO zSbJ<|l6JjV$i=fB)!s`%t}P_f`EtNi?v;RP-Mwm<&wuxyYF`Jowdv=_YeCPhg1r$i zLawX7ji=rl9UFyv#!jN=E$GoYd9%=)ofGJKv)(wz;u(UY#QxjC-d2LW6EMy3yM}p; zihTggZByO*pf#`W1-;~P{tLU4YX98s82fCr20n(KZOLxeN1(OtJ~TT%)?{;Ck6;{h zt^cI($@_FNJ?z8Z#T~%sr$LX7?PrBQ{4GP}{XFPV-~G1@v)um=a%07xFF~u`FABY* zm!V+$cM_hjLwiol9l5>^dcO*K)#f(=qrB_=7Wt#iEVp;y`|$RJ?fSORm zggN(J(4)S7Z5t#1J_@U7AQ^54z@N>XC7c})tz;s^y8Ze#5zXgoCT<_0) zNljeubshUX$f>G3;)2(vg3}rG9oS~jT>|6u4H^ox4FmW8zUZ3P$v&$_S?$^6f(5_B}FQ*rC>n&iq zc6JVWbX>XwY`h!CXV-uca=+ty0>%C);*H<8Rc?}CPoKdi4VcPhXJ)&flNEB42TY$S zrwEwZyR|#1e5{WSwXAY~T3&qpQS3m-=`-$>X5a4v%FF#IE&PLBQ-fE1pDOsS*fas7 z3eQ!xf5fq67ba)EmUC~3t`XCkjV#WSa3AVNRd!8R_~ZLaGCe$hTb~!E4|&C$j)@;X*}o|{rgGg2xo*XJ zl%4FvbJXS{g`R$G3q$^_ z%RM>jYftFWTrOJZjc0A^xc7pb>RrtAc)m1Fp0%wudmC?7f84{P>)aCHZJ{!|77up& zn69i&=Ei%KM)~~^cC^m>20fG<&Hvz90RA`9^B^$wcS+Ne<$?Qw^gYc|!KS{mS=z8J zG;Npr!~H!?{zrT+2VPyPmIbXjTL!egkMi%VnQra}(!8t`bgu|{b+NqyXuaPq54usV zcs94{UfFoFwS)VK))YIdf>+n|RZP~$ovfZ8?k{TO1NZJ|p8FO4a_>$nRoT^4@bbI9 z=6E&Z&0;$Udi%Rld^Uqu=f~Wjg2>3bNEh9c{c^G*6JpOzdTpFk-t1g zTi5CT20gldZWb_V^||v1#zJjqE;omqj?;*sN7vGk0aF_bSsXql`W|akA;zu z(S_U=0n>5a(y*+SdSI-O)_IyD)@OO_50tR0aN?^iu?2-#rd)1f*8}(myAFmP)q9ZX&1B!f znCN^u6ud322%m=ppN0r_m| zcs_jJMOGu+TcLB~McCAJCj*}+6mov9 zc~%3zS$$;r;NA|6|0LMZd2wQ}sd1egFntC(#W1fsl{*cXj`OL3mm&LDa$ktv6HX5{ z6+5HQ!#yPG!08H1E3j?p_g!@5M_9El;8dmJ$0=pz&Ei_`+ zrG@=JU|-ku%L-neA+LG9JlIrUuQ05WJAqHHktv`@{ksahs`tvESLf!{0n_he*91&! z{aV9(F48UchiBX??{&tTJ+Jc2d*!_myc+inpmhzn-gNnRDt1%A^uBm=z|>#v!{I#2 z{Nvsl_3t+5(eIbHg4X$POJRq5ZyL?zA7O`BmQU`r(Y54uvyt)g-_x}g?kN134ZMx~ z@pJbx8{EsHIk~g&hkIExKX(~#mTT@;(fjDV;MI4~_kh;AzPqr?{VnR_ea4&3AD*qR zbM^u7>KwV>WWDB?HLvv%#Tv}LGHofQ0Ir9E-Um%@mw({4|GM6R5Y*a!47?iqqoDQs z$RnUN&kvjKNfO;Wn_qMD25jk^dm6I34m<@~*WV{W>o`3DTIa;$pjlU0TyutTc^sinm_cNwD8w>w!fDGeaAo`Sz|5%KZ#`7X<==i^2HnW=DE5!3kkbN1n-V)dxdn&-h>{_+Z&)YZ?A*a7+wRdWArNMM)~00BYjqQ8@!sIw+dPA zHB!Ic1+V(?j>&rbTC?v3Oy%AWn96++Fh=Hi+uY`-+f%F;?o*A{Ki~V5t$)iQM$N~E zh5aV*sttXI_fe2*A*)><2aJFAn6f=ejvcrRIrBP?dv-pB4ZW{?QrKK2aM#P^FjMjZ?6KLysQz&aOLmjdf**sNAtdjFUMnAY#?pmp7z4YWQ3 z%nDlP%9eAtf6kl*c&pjr$IPJh`FkeNdVd-PdA-NX2z;Q(&j4E2*y%wJ5_~$)`fNBY zXnhAd4QQQ@Q-jv~$yA_qu1^VCzmHAPncB>+`?eiPyA}+-bq{(prwbY8dUS6005-;z;&b7khmiZbZHQxQ zkz2&%^uFS9eIYkqczc3(7ois|Y%UhCtwgR@z;+d^cffRB^a&Wn+*kf4vF3aSl6EZt zIjx<=O^@f5Y_5lA7wTHjx3EdxIFDS{wOsgS$6 zZQ@V7$LqalwL)&GK6GrNoZbUh4|2^CL$hH%KlEPQ3QX^PErFM6JWq?^-FuVh?H_FF zJ#j$5)aDui)A1h|Fukv@88EHmK><_0ZwUET%=a63y(@OWVx%VKTYXt8*jt^fcCBrg z=bL}_93BFl8t*zmkMjCngG|r4m{utXKHp`xdhN_+Es}zxxqqKk;RQ!p6to)jVsRZ5ZSf+o-@cHZ02r_d94lHU;lE z@pY3z_8ZiT`mkA${cj=rEo60l+Z?<)H%FLkmg{T}vDZnrkUxgvnC}C~bdDDZ2f-vnVzbGA3y<670`W@i~U$KFT1tIe%pL+@2vfmVOF zG#}hw%^&wBsK473{`%g8EY~lZzga%{dm*}hj1B&cF`Jnz_bsS@+ksbWdD}vkzb~Tm zV22>Py~$?raPI==tmmv*#$hMOXAIj8_cS!Z7AWfa!2(;c z!1x;`jqGucg64vh#}M}z>|Dg_zd7Rh*5{J(CYRL!_YXADs9n2)S8HGwll8G8oBPN8 z1NvQex5B@>VTW;c=7(Lb=S1X3>u^u#X>lfe?g3iu>|Xf9eG6NQ>|Vy(MgHb?FNFUZ z`+nfn=fZtKYaaIjt#f~G)1Ae15MokY2ZC2)I3VcS-*jbp=3WL}gAOuYA8Ur@F>!B$ z)|!sx!9h-AKO|tPGwx4NjK5i;_7vkDghsVJCB{j|{Lmu)jq!a+b9Y$bYuwAA&r64! zTs9B*o+x~m z)wS|ill7X>*p3evUG;kUYAI3^zYpj;xDyJw#X=vJejh(E$myIqDPZc$7}PB;6e_eJQv_jK^;dT| zeLE-krthlH4Vc=v#hwd%40JAD0J+u4Y}fgqbv(}tdexr`4atLwS-^B|TplpZ)fEAwO0SQziyBz22RYmSd#lTF zA4U^lyRL!__2JIoqSi0r#J1z1(0nvK-6|J$nDR(RjPK zgZT8k^S$)y^7mW8tKSQ50j=*LZU(LOf0OCX;)%c4<2GsTwnFZ+qE52k@^QV{UK_{W zMqTI_-EQ`?{Bd7Si{#{vf_EgwMe}v1@n&_yeLA|X+znnG)4NPI%cbw*$@HFs*wo*9 zp+|eH?kV*0cYqqr$$F~Ia>HuzgYjqJ}7?CW>nS71;31zs-fKWX0^+`f*_t3gih zPp<`xo_hZMw}Z0WJO{sZzP(=Pc`5j&HkzzP-EXb!H%zD7A5B5K-UOy~G-MtEd9C~e zUzPW*!q>M8?45wE?#|%z-GJ%Z^j^Ry<~g3K_&2Pn3v9t9s3ouI_1Y#zz6Nve*AR-@ z^#Od-+`J!rRo)K+Mim}M_M1MAabMW~@pTqpR+Y~mXLm`v#YPbUB^5->y9?~nh?K3A z(x4LJ&khXi#>DRKRQdX;?ZI$LT8&~=a7k=Ndisz^9ec;b;d9HYk zY(ebww~)&9UG{vI$Mfg2JRauqvb3q3&-nXAXm>}jo2t;ZrK}CvxPC?Ku-?mvE}nNw zn{77E^qDW5la@8zCLUP!|Fv0kd&BV+MHl0)ATpk(%ZrSCyPWyXYNpSCm8+RP$E7UG zUw`8W$FtWd%RfGMwW{P(tFCQZMP$6jt{n1TiR7=(Rh7%%eh2cE68p-<|^V3*E+2;MN&pe@@HN;PyE6Kj4$T-ecCvo(dCT&YLe)XB9TJXD; z_`yD1Gx)`r=h|zPtbT`v>p!+-?ch_xd%LselAP(YN9Fp&pEb(*2GS01a zBI7+wt;l#U)JNn25Wly``1-1s$T)}Ah&&kdt3}3ZcR^&FV;+`y6<-H#CmiR8ZAHfS zeYX)A?^o{^dwjobYvK4_*;XRsJT(Dw8ywhlU_$~^L2*KEuu}u~`K&RY%O`=) z;o=ix3=1)4fbNLECITB7*a5&s8T0sxJU9Jr(qz!xG5A4U{Z1I(UyKL4UBnLa*;!;9 zQ#%E}_`J`qfnm%h;a&jej@^je?t#sOT=xhJ+p#BMqXU}@KF0)xZRi`?fb2)Y;(PR1 z-?0`i>r-7P%==ekjSIQqcUZ;;hPl1kUZs3Yzqn;7!`${FK7E}pyqEM`eSI@$c+T!k ze10EZ|L?1**)}2A&2k|7iGeBdT>M!f&jr~eV&}hCQl$p6;ki9k>@c?}BI7;Z4ggMEnC*BQYM&x3u8W#da6-`y!sFW;rwX z!E<(2V5;x7J;L0J&;88~cIb0XVCb`1cs3eY&kzsUSgW`9KeC%;^1}Oxxe_0**ZYTj z>eSe_1B_*UkCK=JT~YQ2iZ0$){4e;e1>J)JL%V|m!+VoMjAi3PpBKXMaVXK%=ZEk! zw8KJ7^mBM%+6*59`nx7L&XJ9k7|OC-^?9LQifucB#E8!isfouQLSpFemSF9UBr#44 z_X_Wxti;ZCo7gjnuWXM~r1P)@P*fKKmGppXur| zQkd(pL|31g8f@*;-EpK{*>Aac?*2?v75H*?uf)LX(D4!j^P5j%>hHjm8!0WbAkEFYWDyctJ!l~UYmJV^UPAEoFeu2n1e$OYl$KTb&dM+e(D@$MKc**7=eTHi&^vgxWk3O?i?i|r) zvXBi|(YA}tXA?gld#{f^kA;1Gx#(iME)yBgr%Oe~dR=0^y&$OfSA8Z7 zb#E43%=spfG3Oga#++|3-z^}X=d90^Rm1VSMHkQYyF~5-{7#Yae7!?tyq?`IGWPLp zB4aMMii~}8i^w|GvcA!0+Az2Kg0FjnuY1f_RzrOT4zK4Ai7wXcL6I@<2Smp4cfa}0 z^3-SMux}q0UF@&L!Pg@5m5r0V!#H^~_vl{4giW;{m*n3=k@2?je|5*IsdHqp{^P$K%-+Ul4z90R*$Qmol)@Lc^Z z*kOBr4-Dh$vy|m>)8Ch>R&3iJ<|jMn=S$r%=D(tg*Rp>^#y}mbTRLyBICK#L}VPhjYY=!x=Lic z)>VpJuAce~Db}F5>1JcCg|V#G`ur%iVL8#oHY_W0A8)wqmk}BJVrh}_eSxJ!E|;@D zZ;HN_C%XFlDaKjBbhGy9^QYxv>hFEwJh>wA%~uuJjJ3b@ zSzN5eTB3_>U(;-}JX(Z0t}D7Yj@JY3FpQ8P=pV$TU6@q=Rq>#!;D*^7)deHOUfJhz$ntW|7V zHxi>Q8ME1Q!eh;+J|~RhWpjyv(ZzFXTVkuv0%HxTMHkyy zAh!B!FxvJKUAzaWG25(9$4bt4&h|Flte*O;aJjMjO1IMcw?5`4nz@lnu^)`z0iuiNbbqtWa@OaLu}*_Uw>S72Br=Zof#xf# z;}3Gp!+AxevHG8E{+x1_|Dq11b=zKI;648k5>uaJF4uPpr7l?iq2?#6w?4;=eYXSA z)#sYA-os2ctFJ!CjDAOoF6J~sWbCWq<}0hO|JGttZL)1%Uwv*FO-=5M6BB?jmDN zcQapAUNqTzPWrqxj{ni3i*?+S*y{7vs5@44u?=I)HtS1$9$V|1<)qJBWBbO7A8g+^ z;@6*{&SLs=(O!QX!+VQQ91nYu82$`(7DJz*E;nxcxoD5Cn%)okTr}RlO$hO^KPMW? zYWGDr-zJB4O)^{00b@=ThIO7|x>-9f3GKw|?6hEqV`6$>c#h66=DA^g`b&NAHPAk$ z>ovgZ^Gsp*+1q|1 zSXQgo!u4b>@$)ljQ?pJQtl=IS^FAQ>#JTf-fvFeI`(e2+QJ-1fy=6YZbMQdo=Qw%K z3UkA8bdcG349xowVK|NsHr*<9CmZ@XTo|_NFp+VbA8NMQIQqoS5BICSy$1SwuGMx@j-4#2)VCd9z` z{a9m}AN@`z_Q7$c+f0Mm=Iz(tu*Eu_EV?)^oFp=~`$Um(zB@r=ygxl&WSlqVi>$TC z`g9fPZ|u`kMYoqHDEm{)wh0{by!CsssC%a9VtvjKxfbkB7a7}in)%N1((l<~{mvF$ zyq270w%HhdTjmk$dyKh&*bT93d)A+;Nu985=a?VY?E~#RPZ-_@pBr?sCKm|9bN77H zZ4Car-tl*ZFABB`MaDVeLi6Q0s!R8EgN#q?yNgZN>xuba5*Yf)*2VA`+1eM*eC>-Y zH(mPz>v^fguS#n2tF#0A<#N%j1HH@4Hmj$ucOl*S-*%Pg;<deKGO>q1P_y}?-4 zZ!bz)`e{nG-AHoK-3bY=cZwgZ!yP1t&7||O_ID9oT~h?-{ku&!>*L?$9KyQXE4p}1xW{aL?BF?c zAIV{^o^F~D>!l4>~ql=#7O_DPYk z-A{;&=g8wC?+5meiCk{1>)Iy0K=&EZ#kN09Y<2CE8nAsX*gk8vJ|7yyz_L}Ho&aZ|zs^)Fg zHCT$>ZMLrmU2N-{!mw}O2)gL^0LcmOsoo;G>&lqIbK`BItLwYqn0Uu@vwrmTUcA1T z-KGy zq1k5bJ5$CU&dDE}ZdRvHNDkk~Gap#*&qNp7^J(yl`F>9P{wVq3*!og*aUS`C*lr>H zf#=oNqKkF;%51a#X(^q9{qvpZqOWg7#&h%=^Oe=@FUdpa31`^$?@11uO9EKapF$2l zij3#k55X74`ArzU_Wf04tm7{tV{Lv8zVICQE7<-iGUo9|@b$a-%Eq5Q6Nh#E+jPBF zS~(x*n}qjJkolTKS&S}nu3?OSNQ~X=o+|Tmi;P7a6aSi@CYp$Co~y2Bgu0c5((}2I z$oTrWLga~!U=z9A^-kA6ssq)=qKoaT5*f#vN}iXlhomeU^Sb^K=G?4MI{uo9jAOou z$m%Prm9CeBy3L91>yi`3Z9#N({UjWVOPOw!Wa68*wIDg*9J)-Qbc`=8a=H9jP zTGn*4I6C$)|K&v&&!y$eHXG9mtUofpx*ikezoPiT`{fmg-(Mtu90N5f+O|?b|H$8a z(RH4%Evtww)_Y~M&H7x|fx>pLN_2HyD7@~kX1ZP@96NREH!}jjCD}S$HzKWdyk2J*DR=JKDKX>0@W$DZLNa-k?ZPORan2ZO}B}f*yj1@+ElfE zQ1-0~rEObJWUSe`B4b1{;DT5TmOE1^wgDLt5z0Hm}{fGMI5Z_b-&*=l2sC+t*iQyk69ajANx%WISj4 zh&;Yg!`55m{TjqZy+p=ZV=hC4VXX&?jQu!B&~I>7}3Rfe6+}Fo3(dSX)oH2CAxhi z4%TLz>1H_{9NM>+=%Vd-5+|;ew>Qz%weqlz6HGU&`R{E?=Zndri}S%Gk#X!#Byn|} zJ&Zer=&sbY^!{in(VZA_oMyUyzBQ7~$KL`wr?UF%T7Q_|4DnN|#Mmmf z$Kt-Ci|6@1W}DSm*Ac|HvqTqt%@i5ull_RVEhLYcMh)BkqT8!MY&1t?^fg;#ZBtgO ze`K8CIe>l+D3pGuW3I^P2k$BWXRN8J*yjCltHeORGv)b{-mWD3LkgvRaInaAzz-66 zD98tjTrTIsgyFs0p{DC?AEBnU9bV8s^7DVa9;N5N5vJ>Y@SHs|uuT}0T$gk3_|rN&loRKHc*UUzpEO6S-+M8QYp&(~Xo_iIpZbhF6)LB2`k0U+NfGR|{1h>Yj^^&$@j{p&;? z0`j#YZx8Y{BI7mwYLRiixJqQ4FRm0B=h7=g#_Pf5BI7)GnaDdeYS=E2_T${3lFuoF zdX#5R%ABco**)7DT}o?ySE0P0HAS~v{4cwfG8`9o7fP8*o_Dq;rn78KOlM!%{ff5< z$HBc45Bu>Rk;~yGbOo*kEr6Sw#&x#ls;X^^3i?MwE_jYD zHdd)BvW>3a75{&v(5Mmmd01p@|9BN`d#q5p&gP>cW6X&v+V*%s|7gfZC2wQ<_NC1G zQ;$ATC}o)IlY#B6spNyX^|L>Ja8;WN#6>_qJ5ch8m#VXj1c57Ozqb zZ1Y?$knvG0ksW`#pno*P!0Y8ZfvKkZImF^;F?6o1QDxhf6!edV7%I7+uk1SF;~dAp zy9L$E?UvfQJlisybJ*VZf=`?)-w#Y9c%GZIQz_4-ld(!wvCZw~$+ciVKPdYT3M%A& zbUn**OxLe0$9zr7EdE%r!+WF;3#EQ)RJ84*g8q@uqq-&~_RYtp>;0yd-sZTrB(lF{ z++aSR1V6Q^Y}=;={iC5)*jJwgrkWmK*LEz|4qcxS+wpm!UN>upuE~h;Rq_~dT}5PZ z?L@5c7X{U9XcOA`8iX;JGGV@^U{+&aC$LIWvd!D8>jYN&LD~OMC_Nv( z7r9)mD#RB1=SQMDOzKgqu5A0Mpno*f4#(2ZfvKkFtLp@!9kLBcT^}Y)*Uu~0MqR%S zU%UKLsMpPEsq52WEq@c;?N!~jUkmz2L!0aDxbA*0=>Pec)b-?WuKR=N>iThg)Rk?2 z7W9vXm@0WKbv?OqOxKSqmyfOwSEJarzY6+C9-}d-yROexE{3kZhPC{=P+nKpTf_U5 ze~9iSqKozUm*^fD>Rz8vexJQ#$hT6G$9wyVVq0m-I~&%abS`TW{0st{rpB^yt!tPK1KsALi|1=Iv(0kQHO;UsEkX`O(R?pO{OB5H zn8VVho8_x(m|@IiMHk=KS%$<}CbV@q)2)E^dwq5NvvT8pSsUAKhq*5w;@5)j6@_8H ztYEtCM^!z3w!WFOrZz6J`smtcMYk4rWf8DmD_Kmp#W{X8VR*l`s>nDNRxw*23+T6{ zFl^iErdt8+_x#r&Y$$pCq-&C4pWytvrup=Csiwz^>wc{*x_BS3mf3o{&{o&&!WykZ zbmQ7w>zZzbI<(DmT4H^eFaXW=z#JrZNzFizF0$M?0fviQ1o7S*(M7+lgI^ryPZTsj zem(T{qbe*1-Ff@FO3ZS0z6Mlgr)xl|UDgg=gQ-@rZ5v76c&%(>?WlAG*?T^^<`TB6 zt?6btTuS;bTT{tn^s?Nu_`0SM#&1V*)pd;U9;toE72CWq@uTY$?F6;h#B>|OG0$Ds zDe4EuJBcpNZyiO(xE;(_1=zT+xHgfhINOEv@7^-swE}-#Bu4QRvtyk}OkJO-+<4RV zhm>Xg64w{%YB4h1u@a+*)tzjc5?x(S2+y(2h_0>~gu2~Kx1v5w@qe$Mt_g(aM^V-E zdWg{1{?^^*BI7lqdr`UJTpj|tTNumQ-9g4~9q9H9eiV6rzWz@W(08Wm{}j_VwvG6~ zb7E_eaqitpWE^u_hPXJVZ5tTo zd$10wQ?{<|KLZJUH5UxPjj!(oo&vx zA??z&X)wk>iGlYe14ztn<~PgzB$M0U^w0elXx$Xc=PPu zG;5Qt12bCFv293*f#YoZz|_S1V<=%e5Vrqz>ZSDijP6s{bHVnH5I;DN3>O*m8%E+Z z7hmOS?rWcTeQ@lI5}!EkN0J!N6iQ=Yt~;8Y$574x+S#MDT|1d>l`7fhy1GUJj^$lM z7tg(&Lw-0X?;4o8cE2sj`1zrmYBuyy(TyIGZBG&K+O&tr7;|@#v5$5u8t3P_T}o+M11yBY}-WQ zCwqq8+oaFX*SHhePa=N&S@;UjcR%_pe7Rip+4Z7Zvzu%&Dic5YtUBsWG2MzpSDz^_ z*H`)+csZuec_YJnhN&SQ<}%G#mbX6RjdSI6qT43St22nMKI4t&?>?qmp-qqt$MwE} z)zzyM|KHD8HmAjBx@U=QQQhok2ETgP`#t;JN@x09Gu9dHUJd<&ZJZtA;XT|OV_6^Q zv&7g3`v*T5KYIq&^VVli@%7wX;%D2?4veACnri#An(MQpnCt%}2KMCvq@DUKX|1}p z?ZA*9&I<<-Kl=P9=5TQEqwc*XW60d6&xV%kXMMi29P4WL1X*9kXEYB9x#2l`XkZ%A z^S&dDLA<6NM(p&NM$J3RO`l=J>&xNhC%ZQ3Gl{5sl<4C9!;vCS)KIn^VZO4s`ix<% zE6RQz(beY!aSR_#boE(190$jkZWD;*^^DKz9WT0dpf_J+d|iE<$T+T#HQ!nL^jSST z*H19r3h@759-@oq7LJFL#SZD=`aBuFX1&DxWY>iF{McopJ6@gHc4^1~ z^SwMUHSzrPc`(hdLU*<~)91f1cl3Fs_^fjS*Yl`O=-&>mLb%{d=?d@j76;Zxe=d=&d5- z`Fjh=;V*k_P+>XX`Ea}0Wijs()(7&wQ)I0D9cEjR_`Ta$*6vrOzC~`}rVBkXa)o*?W=!k!}R zX~Lc%>{-H|BkXy?ULfp6!d@cm<-qVd`iilZHa?4@{r$NJIH!CjGM@jRij3FuPejIZ z^J9_m-26ynJdZvU8Se`}5E<`7-WR#Q<)XXyL>>V0yCUOxxI|>UpLj=PysuuVb?JS< z+rshs^p?m&!T*~ge9ms!r9AsWVw9Wf^|>IthJ0zd9$zgz#*W?V{bcOpFqflU-%ZQg`SV5DeB4}e z!`ggBa*NOVXvER_pKMejaJLvBUU*bjkWAO2`8 zi=p3<#n-$)1wW%bxa@y6)(nn0*WcmxzZL%>y5m6acabs9Zz5yceieBi(En3pJP&^n zxmpgI+G2@YfE9+x2p$T)%W;EcWC%qkFODQ(eGNqaO|yV zy544FYX?8rhwBQ%K3pg0Vt)EPEzDuP5EEIe;0MRW2Ewqm>jyuRG^lN@1JlD^FaPdP zHb(TjL$0SbZNv}u*@nSyacHWdY#RlJxwj2W^*xUM&Xe=F}f9{@*V6 zaUTjdHs*fxm~;JYEB{-MZxa0Ms0VH95SSiz+Xdu2xI5fexsBR)G@n_`caYq0UKt*8 z?L_j~p<2B++>>C;&cSZ7s@m2iFsymkz|iNWf#DpxSztI;y9I`0qkCYQqQ@K_p7HhL zsksTXJAt`kZoHAhUdjrfvJh-dby473iIQ~Rot4`HIsY1a(%l^ zutT5Q2Buz|$KOQQSD|g&1v{MY3W2GK`}DQ{E6fkpx;oh5n5YR1ea>#9Qp22pY_7~v zSo>bVkLRT}y#vGfy-#2mLqFff+9UJNvn%zWZ5~rUU&edYS`tG)GcFgy*8%sMs;0-# zb->HjMnC7>OR;TrA^tE2vhQmwYlp@tV$<&z{ODoNCqDa*=XU>K=Qe6TfUtpqVZ0UP zbtu*ieGVdax?V5ZspPi0J}=f3>oi1sj!|seU=qVWORjKb*?SEA9JgG2{T#I%n;ZJO zT2*ZuD!CNXF?sut8=faS5H>6@oY#j3HpGL=eneo3yuF7e??E`zb@R%#H+~L+^ZG~< zZ@X}>j&seZV2AU-j)7qt1_qzVCT&*VhDzPp=I!e)@Ba+7#nH1W;M85VpK!yaS~$?$YV)<`piG(H=g8|J=gE;Ton3G!#kT} z=WYe$UwVIc1o@(N7Hd66bbEvCY>_e6ERpdXo@u@-!N%*@TpEO*S?q7R6^X8|gW&B| z75Ag-AYg5Jwk~b&T=9!__4jVGIJ!3iQeKJdIcR_yTJ z<(S|{Exn%EdJoP%l77<|9v|&CwEoW6!^ZqPu2w_aPLN!%UdM~v$D!;iM8-awFEYN@ zdz{FclaF5^<%!;tOkek?A%`baj1+A~trJ2pIQL zk+JTVh>X{xi;KqjoTlqhV4q(px;SpH5E)-*U2eX7+-geRCS9Kb`}=Cq#c_L;*;XcP z)HNtj_ZkwXPR0Y~e4XgxTzak9R#+Syqt^$9`7g286_sXJ2YKEoc6c7$5d7fhfHxU) zKNv&Tg1|9ybMS+@M@v0%Zp8SviXG0ew-7(R{zI0JuK$4f+(!Hiw|u;=nCp0HSGjY0 z<{Ok#ld!w*?Rt%hp!!x`PDwgx;-L(uucz?ob+=-oFg9-U2N~8W?QMr*yeR> zE^({ zo)=x5$DcFXX5i1;pr4zH`G4E1qKm#>5gFU`vdB1AUJ`l2f7Xo}y(luyO)rS7ZOZc2 z^*S)t>!OSO^qSdLn17rD-v|tId^0e-9=>HP%crf>0{iD3(Z#mBZMK!rUa#YmVcaY+ zUC+HJn%%p^&v>~GVO!oa-7H`I%o+Rc1JT7Ay>GUj1D>ZJ2BxWbTRsX5^Z7V1v|D)# zmGam5-%EY)dsLs8pKN~8bybR3-O2X32$;iXBIEh}srmADs*Shh3uBdf&^BisBo6xh zlK9EiWN|+&Hdn_DuQ7^lEY?@SXPq6_-PgvvZ_(~M5%3)TR%EQ-H)dN2{yd+!p3C>9 z+a%G|?*(8DeimJ&cs~7LzP$|^!t>DeUU1C)D!Mp#{9?8hNnTgC(Sr>& zRmuJKl71`4bUhelo=-7^9+K@35ip8A&5`z!eR(`>yhIFJ2p ztkP`pn(%M%^N;CfIp|t4c)nNkD($=aj?B-$)g*t6TPe0!!&=KtcU6JuF|UKKGvn<< zrt8e8Pj4qO{k{k0-q>Pf=f8gcqgV!YAX_uh#eQolGM)=f%vaNT^^!Qt3ahr`x?4tM zti#eGO9o2}=es&1>_bHRG7F1lFj)kMZTS2bT*pN^Mv6l>MebSoe?kE?6l zVBJ)5=4;zzc8la1S8hz{`ZHLwHAwu0#m>lI2kZB8idaL*wzl}e`mQB1p8IQBoT_^D z;{Vtt+~C53!Efdf!xoj}3(3IA336?AKPo*Lvp5`v%Y94Ta%2Zyj_stmk~J z%q!TBZA{nu1O09k7`Cmgu_|j*QMA7|76JXV7a7~yuBe>P_qz5Cj-5L5sk=?guF8&K zJv#}*e(NYQj+YK*>wSUwZ7K}MKv$8mMqPrh&gRQ&rs;TnpSCfav25)bw~K4fR9hQ$ zx0%KBzSm=JSJA2TJ(6xDw=1OYv9{=Ec5xu*^Z8I=SdZ=`pK9qpJSR3M`Ruw;eLh*c zbp0DVA9@f!w^_TY)SYc!AN{Tyj=e3!502wuq}{W7mgdlt2H-(!|__t9SCwCvuy&$-1c1=OK96) zbc@X~-F_lt|MnFb^Qtr7RiUY~aqU7Qa_nXUH) z=DnLR%xl+RyNk$}=guPIxxbUh<;LMgo7S^z9FCF`1IPRB!4KxRhcWk46s-=UMZmW1 z8T_c)Xw!{rUE#SohS-sy8T0uEbJg{y@cf-9x>(lTRc<7E zXZoF496wVe2G)8qiTR?{Gs`Di`^sbL+E;j<^{LUs8s<)wT-VpW@)&grZJSDR`@_as zR@Y%Np0H1*5kE(S@04O~rkh=rCnS5%Ti47gVq^P=fb+u)vu&E##x=BNi!S<_B?6AI znIdC7_Y)cCnthALd410y=Y)Qr8FQXveyWn3zZcy;aD1-lVov)LTV3m`h|T_gBH*|` zz-*go#*;}Zj5`l_^DQP z+YSqH4<&Kd4c{?V7oLNzHHP!j5y20hyGI(!##CG*>{!voe2)vpur*njgx#&hN< zk;|R4ZAg3cd+697$C)2LPsgiU+jfxxPjf-`^GVD}+p1zd#y=tVr{BnHhjMn`8uQ## z)B9y-$*q`<9+d6)&<;K3y1I@U*7}6d4!l?LHfAw(%{BCMV(^18TNVdrUXO$AGl2>< zu+4M5M8=!eqsoEoo&6Q!pDgiljGPqm!8z)bz|_6_{agH&>jPaAt{gknd}iloi;kt| z=NY1l`J65?&gG|t{IE~X3`{ff+?SR(snTPZ8L(L=JGWj?chUQT@0)vFc%KihP@ zoi)JD6Nc?NS7dD0IU?g+x*+&cSKby~3$KXH|Aivpxq3nHb-ww^+OVFqz1Vf;YoX~@ znXEgHtItLhv9XIqs0F=?f^9Xh%Y@;)b*aeM-b;e5ntEJ)?qdM>x?FUz-*>mMR|Pt5 ztIvgC-LDLBuduikW{WvrWz73iRow3a=>we8uN7UK|F01l=i#f(SJnqNN&Yxit~1@L zq&;g$TS^lV1+= z?46=p2RYp#azBu7H(ynD4D-52SP>h$TLiqPx~m8q=2z^qdyQrNySH}|U)21o%nq5Zw}?5tAbxN>KW~1s_N*Pw$(KbJ`{gB(G5vvnRSPgS0vS+%ebm$>j-SzoDoM(TK7&sQb7a6bL--(Rv`Br2c zci)JNIe#s3xq0Yd>2U1hABnE6rC7x3PPU(juCAqs^YG86>orw-ubZx=SS~+Z9}#Qw z3-ROYC1y3&XC-m2|4aN7U1@09{uBYvsXs)nt4GEEe-|0AXTOPz*Ueu=);47AyV#!T zuL6JG_B*6bnAbnz2mAK#;1@q9{5LSPt1nu9uIW01MN`YCN(zj&l_F1pTq}wlkY5LS z$e5@G-6o=oW3jQ>dR?$)O^vkxd#}sb;CDIE#hjNF8T(@yk?|Z^T4c;)DUoq5Y9TV_ z(p+S`#xxUI`!kzg^ch)f^Gc$NxveNN)^i1s@qAo9_{MX6bz#`9)kMa;R}H>a5qUqT z-^wDF8;9+Sg>D$5oxrYTutT;6VQUh$7GY}>whm$I61E;;tq5D6unh=nP1uHnwIOUH z!rBtnjuhWd8<%*F^%jQnWG|8NoU0MJXzsQm;Cx;XxfaIz zb|T~azOBeOw{9aco|jvTjQ2KMiHxtYw-g!gS9*$!=kXRI8Ji2o z`K!CgI1h9a8RxIfM8^5+-wjKzzncojIjF11xCT%ck@5BEUt*8YdbE$& zDy!6;ZO)&Rd0hVex7C_%#eYW{?gO4GD&~IH>{9Bdj`-0vCClZb>qKHceTkp7B__66 zC2v=@rlhmR;jqp`>C3!a!hU1|>@x#`ld{VAm<=XV8jKRoe=^@z$gm~5Uhl~FY z3~V~EL4hf9pYyH1aV*R<+rf#QU*~2wRi%dOm`W4%KiT$7*ar2e(cT`PngH9%HA#@Lr*t)vC(;=+66QXR$-wT|_RPVy3%m@Ov%z*ex(ct&HzZ*dE5R zd`1^p!)Gyjij4Vk|2`PxP*Dl$0y9&u@_-`6E-1XK9(jX%=>s!ku_?x zwAE2}p37um#oI+=Q%rWQhnF&@JLM|3Qo*)m{SVnISGg4&!@rfQjK>y;S_4N7#W0n`?OWmg5J$2_;I3A-j?9{*ax zu1nZL4JzC93G-ZTNSMdFk+2h!`P5&Z+(hhdChQi%ZVe32$J+uMQ{U|3|F;LG2fgig z5OybFcO`6LF%;s zJd`khO|*!x#e_YaFpu|0VA$`E2Bvv>E{_rRc)~p16NEjPFpu|C!rbT6f%$7r&G(tW zuno^9%zZvb*zZ6aDq-%kal+hZ6T+Gj){L;`fsMCX>aIm#ddzcKim;^<=IvW1VeWI;gt^^vge{*i z_qhULD-yO+U}LoLwyhkP9`<;vB+T2eYQo%ZHNsXWtYyMH-Wmz>cxw{27GY~A%wJ2b zlQ56BE@A5t){3z83ELoHo@Z;qHcXhep$%ah5!RNlc7(MjY-7SUA*@5fye=IR=56mp zSZBhz7_0P?NA@01uMa*}6?YAOYQSdGz!dp=!&bT?z0DBTM*Wh zuq_kj?b|A0p3ByRZ9~|$30tU*l&!9f{$F}s^IWz|?3@({n~=P>@9!T}CwAVx8p3)d z%zgGItWUx`UM*pDg!N6B=h81>9`!nP-DC}BGg zHY{QOx@dU9ynQ1Q=5`|!=60h9+mW!H66SrebHY5{E`;q$*lr2)c)Js}2Vr|A%<~*= ztkMUU>^)wbqsJtETz723JjS?$`R9J)6Xx$v?UgY1xi?`G2%AXQB*G>q%-f*r`o%fZ z^PG~{Ih&d=e;q%Ku<3-&NSMdlCt)6M--Nl{euT|TnERZSF!wn-VSXOZNtoO1PuN_- z4j}A*37cf+!+{C&@p({S7~C+rBqjwI};gn6Fx z2s@gvV+cDoVcz!R2%DcU&-3_%xz7^<^YdBz1-u4R;<~6&Ju!RY8pBE8!F=3Y&tJDN- z^IR?!hOa#?OZ>R*j}HT z7}n73ZW4y=y)p6Qx!p|IEeZ1&w-R<+VE!6c>wJ4)ddz*^LD-!M^ETX-F!y;kVfPSr zFJbp3%yYS)um=cxkg$gm=HESAMA+hld7cjw_DI4!-lGX~pN}QX?H(uWiG;b&CkcBh zVg7y6rwMx|Fdr*A7M~4Fk9jW75%xS`FA(-(!n_SHCCqbqIbm-13SqA%%zeHV7=GUT zdSFK#y_+z%doN*b_kP0M?gPR;BQ%Lk^1y?rYr%-gqO!rX49 zgt^_y33Izu66SWR61Ey)t0&CkwM>}DTZ6DQ1MB3)l>J(PDRQ4{6SfXv>n6VEDvGZJ7CCtCGuzteaZi9rmUF(Fo-G&KsyEX}1Xl)qYRR1WQC*0>oiJh~y z3G;aE66SX83EP;kO#;I?q(flZB9GTGVV-9v!a65xp*ohWOTye|*Mxa4n-aDeVciJp zPT1yz^#}}K3vUsaw%OazGhyELEfeN;TM@Q(V0cfxjj>8i!Zwe$tuV}cyWq#qVKplR zrpG+D>V%C^1=(r>>s5ca_J2NnR?R{2Y*bgT(RjKrKkH?!~tWs5Mb9SIGe8%By660$U<50rld8H|M{E5VF zR$x1L^0Hr0t3ux1?#)UW-jkgZ?2w&H*m;3z1phw6`3du}zHytRcEY z+AU0&p92>q%=5gMuuBNLG+`d^vcSCEn(yU-=`pvvA~3YOGBCVPxXPH1WqeI^jWBPk z#=1KA!T8q_b{%2Y6Lte(HxhOeVK)8u-@v%wub`K!|rntVT%*yZGSjn z?(>m^x!t1)bGye9=5~)K%SUJPjaK27ytUh9L6RV8+MZNWL8zu`yJ1S9bR9bCVt|zNZolo4r{Jb-mc3mepZiZgk41#yS_ae+JM)$=LkE# zSfGYBpxx`*JZK-XHwb%^u;)WAs_J9!1;SoT*#5;hjT*f~*vkp?c&{YPeZETAYYB6o zuP4lX-WTo(@Y?rAV&^{JOqlz8i?Fu|dxx+k3G-auCG0)I-Y4vXgn9cqcGgaA7@sP6 zE*}!Rj|lsiuul@^@jgwM=lL07=XNRIhR+i_kM~7jIEQ^17}n)0!oE(Jx8a+Fc`n~3 z%|es_7e_7aTle|HCZ%mq$yr5Rc~(hS6=97DYeHDl!0- zu+<4`N!S{Mt(h>d=UNH#dag~_Itg>1>ju_VJKnbS0@K58*UDI>E6d(<*+3ZjSwHdP zx~&P@kgzs{ZIm$2yKTZemv)4;53DV;d*i?qxz9}!=DBo8nA>$EtW(0=XJ^8?B+TP= zHJ0_~X2RM+dp9M1x+TnW>z*)=zjo zn8(|WumWM#gw+t%i?H5=^&u>KJ>%`GC3ba$^$iSP>+~b6e_%LX1`swdu(mMf2L-0c z`*v``yzN5>+n%tYf#G`?I|PQGw+ssm+diDI5rmCQn74hDu}Y|ww|ysJZQ=Nii67V9 zIk0N5+a)k;)2@W=M%eCz?LpX{gpDR_3}IsltL~$%&F_K75xen(?M2w$giRo9B4LvV zo18EoTT>F|<7Fyg(+HcMFpoEbuzeEdx$H~WeuT{=Y-iH%vxwbn!sZaRKVfqTJAknN z5q2P92N8BKVTTabn$+`9Vs}_zcnvx{Fg;U!EFF}bM{A?#Sf zjw5V7VaF480%0cx)>aE@+ev}xVXw={#wuMI*r~$M&nd*uX@s3l*l6+|#u>q`E#!D+ zV2Zp=XAyQbVG9U5hp=-AJCCsQ3A=!>Y#l{!`-Q}AAz>F0b}?a>1cvt>mj;HfOD+ox z`|onXt_Z9xjEO4)L!Va>b~Rzw5SBe3=KXLjvAd42cI14xKG@;+GHx(d1@-V;ZWh)S zj^8A5aZ0yiH<~SeR_3<15OynJw-I(bVRt0V=l44k=6T*l*xiKPlQ55WZ^As@{LoK0 zw(lc$_Y?L2VGk1a5Mhf5TTIx)ggrvoql7(1*yDsfLD-Xl4To|5RA7p{|DGo78N!|o ztS!WQE-;MuJYg>o_99^~5%w}+uMqYsVXqPPI$>`R_9kI(5%xA=?+~_xuy+Z2kFfU% z`+%?y3Hyk!j|uyPuulp5jIhrM`+~493Hyq$uL=8xux|?C#(}kkdA>30sG-bpsm(wO=nVtbHrO)+cNO!deryVPJSZ zwh2u0^l`9J!h9UGC9GY-+-LiQxzCLW+r(Hl4mt|MHgzC=Iwj0=>zpvptqWmY1H*H0 z)4=fal+6OeG185&?u2cgFmGRvgn9e6Agm{0TN1WaVE7*9)`4N$w;^m>!nPxPCkSbxF>5H^soL4*w^Y)Hbq&$dsP z*J&tWI}n!rc8U8OM(lwnxG| zmpv2axr`<(TT999#t^%)gk|52b)Vyi-FU+GB5ZHMCJ;7}ut|hXCTt2}Qwf_!*mS~X z5Vj9t`x3SvVKWJvMc8b@d|lT{y`b6V<3*ou!uQzb5I_49Ha9T52R;kFWy? zJBYA@2|I+aLj%KW*I|KS8xBvH&l^V=%W8j=FpP0z;>YVWkFcWy!_N_p2@LZ-mayXp zn@`yBgq=XxiG-a**vW*QLfENx-{G3nN1%zEl*h0cC3T&7*%C?IG)5BiFOA_Wayp*uZ2)mrHD+s%iu&W5W zny_mKyOyx)2)jNo91}MLrnUBZ-k31&hno`Sb~h7t3t_hsb{k>06Ltq-cM^6NVRsXD z4`KHbb{}E)6ZSyDypJACnD@~`ge@|b%_$EH!|}72_<4k|M+tk3uw}^mt&bDCClcm$ zdNN^Nr>6*eny_aGdzP^02zx#-?1vWu!#VTCgn9d3BJAbB@I91Q0>gRZRl;5)>~+H4 zNSOEEn+fywy+zpDfnodJ2@Kn}gs^uBdyla93HyMs4+;B-u#XA*gs@Ku`;4&93Hu^p z-nU;S%=`8$!oDW#8^XRN>^s7~C+r8pekAOtgn3n6=97DYeHC4!kQ7*oUj&w z;p_IL0>f*^(t+V;2+IV9K9>y)$IEhrEl=1Agsn)}N`$RU*eZeH`LJqW!?c06trnOb z_OY~jU<2Gp_ALWLyEPK#V|UGjdAzj93G+5=Kv-+SHVkYO z)VfVzSnG`lYfD%=!rBwIF=3k|%xl&mVP3P2gmogUGhtl_>q^+Bgl$GxH^RCTwt2$5 zW<3(-HQR!)o`h{l*zU>e70+cWVz)J6+Yq)bVcQW_AZ%>X2CrE)v8y30Ti?QcX1^We ztQYawo3K8F)e=^hFz>&<3G+Pr8LNaj$Jqd3_}Z;M@iUOHL4*w^YzSf76E>8v9S9pn z*l@x|5H^yqQH1SC*iMA)OxP}g;eF7qf#G##H^Q>t8S(M6JF(k?ussPIP1qR1#u7G; zu<^#S^J8yecz*0f{7fKhB4LvVn@rdg!lnj>_k_~|!+Xi;f#DdLLD)Wo?Mv8xgv}&u z7Gbj!=JVv7g!x$BKVfb+H(_pf0Ac?l>_EZ}BJ5zo4hanV_RzqvZx18vaKert>`20n zB5WREM-z4oVaF179AWbVE8c$Cb@ceau%0Inb|PUX5q2_RrzFhB?x_j$v3nX}rxSJt zVP_I{7GY--wt%p62s<}0oXgJ(4EyhV!Y&A`DB75~FfjDFkg$sgyO^*`2+Mwt(EI&T zVs{x~mj{O5@w_51Z2OhQvTN+s!mwYiB7UwR>{`OEBkX#@ZXhiCtzED0jb`WV#rv?E zg%uzDvoUiMiIJ__?J*u}styi4#A9H+b-h-6p70inkv-d~>$hTlw}~#Ut$M53Hnw~* z&H}M5Vq@uMD7iCkBf|LdQ4=@^HGs8mq$b{*WY_eY>f4k z>1MT@9cuZE=;C|BPZL{R;}_fboaka3o;BNSyiO6@a{24}z*y(!%}=G8*yeTC^@6dj zFPLu8ebDVL9fi8DiZ0gp6_N4$ds$@sZq`d8V;{UIvbMq7hig&3M%e3wy+PQUguO-B z+l0MC*b>6tCG0)I-Y4t>!agMIBf>r=>=VL1CG0c8elCu=d`x^!?7j$0-S`;)GGU8~ zaT+!HDq(K-bzu9Mjqbh)Opm$GZwdR3u%nGbtBpT7MRykRjw#b*bRNF$2`y9gB`Lz0&^SH`7?C#)LIVn75&6VA!|K0>id9Pnfr_MZ!GZQVDarr4#0M z%Mi9~!rbR_ge{*ikGDd?+~Y)+xbFP)yWo>F9H(}WR^+OD7(*}WI zn_3gLZ1F~~G&gV4hQSWwwFwO4ZA4gG!rBFf_totK!+6^yZSX$ZnAmMXSO>y764ohU zo=fMzu%2B4!}iTc+UI@LHP{VPL)$hDOb@%=W`uQ1n75&O!rbTP33Iz1gl$1s&xCor zEeYF-u&oK(CNS)yZ3Dx4Zbw*wu9xx3b|rSZ5w?56+~*#F zVcYi%Ofz%4(FyZ5j3I1nU?bJ3ZQ}yd!|ro@VEBIYUV)+A-U;(IOdxDx!aSEr33H#5 z6XteP2%8!h_QSNmw0$0Ldcr*348ryyY|D-sq4fP`_o?qi*C@7aU*aeGEiw0_?_EZ}BJ5zo4k7GN!VV+s@WA#h*0BD+BLc(o`^do5!q4xc0vioJ=LM!{mhqzrJBF}h z2|JFk`3dtj98cH@3G+70t<%PoUJIR_80@gtCk3Wn-0oz;P6-VA=+wZ_=V^hV-RXhh zdtYY+hCa_E?5x0WOq?AU`dmQRIfR`X7`_L2USJsS{J^k}E(i?mvh`WL|1J!6^friX7YVtx4WL$-9Xrlgxy5g&4k@T*sX-!M%e8M^IG4LFt5v%1H*Vv5%zRo*niIihS!H@1DmV?ZF?>-J?wq^JYg>o_F`bz z_Ll;~crOQr*YH;YQ#(KZUM1|cgn9d353E|9+Q#nd-UxPT;y(Y`Sfz&RN{x!Py_wi~ z8{Q)9ZNlCmYzblS2DXzYDf{;VQ{?S?KVja64-)2f9}@NvVIK!J67u{cFwApza((}l z*nLLW=Y)Mh*z|UqS;LsXJiiQf_&t`d0zUQi!SnHZ z;^#+U=;sHb`*YC!$#g3`IoW%DzYssa3B&XESK=r8y$X-KKVzse-K<|56V`;Vri3*M3~SolSauF9C9HVKxBhBjx{Xc78ZRrX6&znqWUSpX zBI9>jmNs8m4VDiK^Iah@^tobSnEOhB4TD@)4h-A13Sp}n%W9{;v4P`kHKN;67}jca z5_1j0)(i~mwwAFh*L8$pJ=P9>upWzLEbIK8`O)9j!1ou|C4SZ;tW{v6G^B0o2Uh&Z z!`KFa;n-*$7>#|c|=yPYn7Kgb7?RH7*yzRRZwp(C0U+f-OjRv%BkHGY> z*K<$8Mkmbk9Fs8jIhL?-f#Eqe-dLq44tee^4Et=a;0N23{Wgf_Hi6hpBy194lL?y= z7+%+=28L~zM%Z-1mM8Bm&j@xn7wr=m)_UK-@SNK(Fl@ujz%bsdz|d|sVRHh*dhQ<> z`kYJH0fAxr{ukI_b!yvzf$3pCKMo2EeI6W`+WGiAgs?*iJB+Xe#Sf<&UPIz*;lqjD z5yrCf@hD;F=SbqGCmECT%&xH(!8WhOal&wpJ2u2TMr6FlJlcF!m>;b9{J`ownCy=a zOp)hs0%0c-b`oJH2ZpaFP6-UJMW+Uad7c&+zMeSUSd|*cR_u=YKO3DTtcXoMQv`es zd4}0~P4PbLY-8E@I7e6!o8JY&j~@0KpBosq`8>kT4-Ch`1%YWiKQ}K-nCG%EFud=* z$XM3iON8NgyEyp4HeE{CWr1OxE)NXjT@e_bn^y*gc2^nmHevr>BMj%)tAj4)cAYRh zSFa7a*!~-Y;dTFd)Acci*NGbg!`yEQ49}yR1H)Wz2@G?&H88Zh&6w8;*&V{L{kI1{ z*rq!JLqB&Bb~jB7v zY>_I+_E5t7yj?`tV#3btqMjRSqmsvanAkl+*q~5v_33#&O6;}{?-`@rZ}xpeKSz{# zo{t5ec#raUVCeG+V_DxnB@FxQ$>0ai!>0qo7|#&)Y+xAgxxmop^9l1>ypS-j{jZy; zm%OiEBz7+m_Hx2Jmsb+zxxAV%w|gyNZufe^-0ls+-b|SLe2cKR37faw|F5U}d?(o9 zoUXbDvEU=D9Qr3|~_<4-DhAAnfRH z-+=wORInSSPHkH{Fg@&fE<@O|ge_+*>(3R0VSg@9{H#dWN`c{h_{xD{n^qxg)xhxD zvRYv1b9KU68p~?2rZB9<8o>{~eq1Xs?BBHsTPHBgdtJgNC(ma2xR^@rHP<6PTLo4W zH6hv74@{4FybTCzP1r%qrA zS;9PiH^RCn%zbW7SdYN)Jl-NO&C_$~nJ|yHWy0KUE5f!83~RPcV3^CcfuY@YfvKJ6 zSuj?q2W|5_YlPu^R~`J|eA_E9b>)71J+}%E0L-76TP9(;ep@RzJ}Bh#su*Do1BK%~ z(*Tk2vyJ{D%kceg>nAe4e(ozW_I;hm!~WBj_pwCA_gedijIRZHi;U+R`WkF3tFu1; zhjE4wU47;c`(*nNQ#0_o>l$#F53=leK)2KN(aPEB`d`TKeYl~POO>kF=05e=JA8dG zOmwk!JA`%=KgKX0!wDNf*hs=g1vc6fmi>-_De}*=cOq=(!1e;4y9B1l&%s>-!(4U? z4DEIg48LQuM__8_ZP=5r(S&8+b#q)hCU|~HYKoq)zG%7f$3pyLm{~jbT%!q z^D#I*VQx2ruzeEdKKD(S``j;KZZ|V9oYQ9ohHamnF!wnpVV=wWgw0Ku`#gZK|0T@h z9hflpc~HXK?qI^=a}qcP4@vAi-k}Nec!wp-?G86qX=4)m;Rs=PjXE;;nF}^Y1%`g+ zCCuX=Z7hp_j4+IUEb(((!o2A?@y@QlE)y=Ml7_he@UhHW~Vumyn)xAy4noWRiM{BX@cb|PWt5})T0 zc79;kFBb%cxm-xtLc%U0>|(;=a~Rm~mk_(`ISfBPE+uxC1y-w#x9#%4^sxK9g0MG8 z|6NJ!t|IJe!mc6gTEea)?0Uj(AneA#aGc$hd|v4FyeZgWt#2mmmW27-duzfx-fanU zyW11yc6TJq?d}Y0M;k}FyDKm~=5}`zc2B}QmwO4jkFfg*dw{Un>UON$A6I53<;9to@{YG=syXkdEGYyDVYcwKosFtzi(eIjA*^GU+ucc(Cyr-`Lgn1iY zC+rQv-V6-e{#Ibv_O}DWHoOxU+ARqT=bm>1!?wSdFmL<&gndBRhlK4<&Uf~`myd!^ ztjouNVO>5U?9+t#nD~sa&lBb~`+~491H(3a6&Uu>*9r4@w}yJ6&u@qw``*d7#O}L< zdCk5j?1zMT8-670C&GRvEc=d#=lM%w=j_*nd7i%ohPD1ZFzmlS0zs`!WghepohPJ_a8wUYPTKSdQ2&PuL2Ct(Y*+bEUxWb@(5^LM8wQ5wN}IsY=SGCJ4Q!|y+SV>GJ?yn^A6ShW$$sO&iZADlZxWau zcAp&x>quB9!m{;tJYHvF*M+dIgl$ULW`uPktUFXL^u9sSheUD)lIh`0?w`5iHvRCwy2z6gRZt3RwQwK{n#vjT|XADA2lR~ezsh$ zpY*fj;Sjr*_`$jB*N*kJS$^4P&0b$$dp65m*Pg9~n7u9cOjp;U#dh{F-C~o>hsV^< zxAC*+{-TR<`iWc|>SothWE`V)B4fYRid=52={mO%E3En*EN+hEhJsuzp@drkH8eXP;;qKjjFh}n9+`1$@&!ge5R7-7Q+ z8$sAe!bTCcBVjucwliV75Z1m}$cFyMdhSZ>b|Y-}!2G)#+R=LirpNsJ-!oy0R6(}U zf#K)aV*fPPA#C5ku;2F!41LZdEc?wZul1~8 zHx1e_J1~qlCt+T*{S)T#<`Om{ymy504j^{_BkVxJ4kGMe!VXE8xBXDU4ojHV0z(u@d@+VoDkT~ z?ojq828KRQBJ5%|ixQDQN1HQEY-Q6*gZhlgM>Xq*doFf6ZSA+qieM#`FRxA%`FZCRhR9N;1j=n}XKBOON3F}=D6ZQpRqewlMgfWk8_>%biim**Qf6S!^nyVAzJG2wR%4WeCgGlk}P`OYHnNocw+UueZyK z9bPY%BYsvQV{28yRvVc`9`)x_s8icJ$ z*jj<%dU|UIhPkX07~Ti1Ys|+1+UXjB#nW65%C=tcqsP3xtq5D6unh?7XX7W!Z6&$R zP0|ExYfb$0C*yEKV%LVS?0dOhi;algFp^8##LmZIyM%eXY~49$?TOvSfw_&Abd$jJ znDY*S;at%%Ft<_rPK0$14C8entSe!g5;l?a!)C;;TVU9}S)DcOhHD$Pp?k1HwmD%v z0>fHw5g7XH85o|&TLy-9TM@Q3VOx>cs}9Wb z(E8K_hBfOI7=D+lcVKGg@%kjpW7ZN@M_6CN`VrPYVV>uJgn6D{kn?CDu^U9#U}Ke9 zEZaQZ5MlneqV0np9Opv=Q+;pK4heI=!w4H5nA@oD5rOG3=OY92zd12yVrO^ezbn`g#N3wI9YqX%8BMH(!F-I#Eb1ReCWdp=v82aT zxDVrq`FYDTo*0g4M-fAv69k(m*wMuBIyZ?J+M7%a_q$CYhH_H{n@sxK)&g;+JF>W)7;k|o%G`y*CzP3*s$DFH@4IH1Y zOzC}GDE+-6`w+#>_KoBI-K)u8Eax?uRR!JdlkFo%pKeHbQTF_LR5ezg!h6+?LWkA$@B~qu;lrG|q)@ zPU&^klkG*vnB0}}VodH#X&i^{$hsnV-(|5{3;pRcAg?)i8n zrPnnF*rw*%tnGi^R;aLq?6TG7VHzjJ{9aU z!9Ex43&EBN_N8E73AS9YuLb*tm^Hz&+zwX=_N`#w5!(^k_4mZk-VXtrlx2Su>?gr~ z7Hp+pzX-M}U>>ty1LkG;O|aF1Z4m4`aGk%C+!%=UAH)o~+@FH|CD`AB{UeyKDdDx- zefd|&t(mB*mJN--g5txz5c6JGQ6gjAPwM$TcR0bDJi_ zFlJ4Op`XnJn&zX?Pb`J7?xot zVkpP&kt~MgY_F4$i}pjiFaJK=btN{y1oP32n4R|e>n_+{f^AQ6Sxcv2J%r8@!FmSFuTQ-K z=H=@x7{9v_>$Z=O+gq@{#8BtF_7@WNrZEJ^(*sBkj#CE`L%$Ca>|nw8U5BXi5R$|CypZbC>r*O*2PevXP98a1sfsQNWo?XdqupTA3<^$!%>2bCWbK_ zBiO@%FJAXY3c0ZXbGdPXjSrZYz1z0tb@{l0WjHF3b9)m4<~k=5!+0J|4DEdy#Km<^ z5^|G?p}i@@(B4$RrU|xhp0`AO@R%fr_Ky{MrVBPw%#EUTaF5d*G5*#K*3w+=IN{6j zf*n*mD306D3?VmDuvvoHn)SNc>ue!6M=<`57}nnjft>e&69eY)JV~&Vi4{pJRys~0 zX6IbzshLIXZf;^Y_RbS}P7`dtV7|T_J-=Ec+QI2U=NXwrWwJHg0TzIKEXa8)`-$|| zV`FQyJ>mSql(*r*!)7zvH@;z<6An(%4q6DKf`(Wb38aZx`phQ9bZ_=`ATQ`gOCA zwe?XP4{l3&v5dE7*(mnBUV3NBTk>D&D^$26rBV0oS(lF=Z6R-u;$->t64_ls?ry>E z5$s;U?i1{O!5$FoLBSpp>|wzk5$sXH9uw?w!JZK8Nx{kldrGjUiCJ?g%WdWv!JZZD zIl-P6?1g~&xbb4Z{2gU43C6$O#_{lFA@@fyD3$jWoR_^KK&3p_O4*>3HClQv+jN`C5COGNr$!lw3OHB2PB8=L%}`@*hCXf$H#(wBG{*b zeJ0rF0rT>G5is|288NJ*wsox}<$cRiF83wL8FPDI3AQ|7uJh}Fxz29_=5n_M>!Z$A zkQ~l2za@t4&G4!*0FgCb7Uz@_YrZfBgS<76H z_UE5u1NHtbZ2n6OW3y(WT;7I|RovH7F{?xjeW^?gbw=;<@G@2j;$&ItD z`PhJ%opxXL7wv3AA-54RY)f^CVfpG2LtpqCQ^@L*9NKFj*l^Kz?+?}qCz{c8Y)m>2 zS#!GhcN1cE(&aW4?5GmU8Q0&Tq8>IAI)_qQLSGsRxy=P@L=0oym>Ak?LJaq~G$n>| z9r{};j>~cpvt~l)76J3Vxn;m6nhWXJO0cboVO*w*cGjHa%#Pb@Ay~AQ;c_j7T&sY2 z*|#BvbBfl)%ukowHeeGiLprt-tc_r81#3qPV?9^Ix;@F6e-rJWbaW7Gd%-#iwu4|h z3O0=9FnBMvlaSji|31r~3t?OqQe394DR%L1C(?<}t#&460u!@rXTiD%whJ-zb5~+$ zFItzHn0?uek_)efDQhVko!f0bM*U?H{aFOw4)@5UeZN!+PjO41MV? z*j|GDO5nGS4s!uFaf8pCciFE^?_az(HC-xJ11|)`MEG0b{kMT6FAlqNaO{6lSZ}-u- zf;yw$=1$CIK0w%8K=x4QfkN&eV(90=#7w8R%R@4Y?$d0IFx{&CcO07XVm}!uWbL^f zUPA|^yx2Yt%d$RR7u_*i8cfViPt45M2&?2YGFv0WIrY$_2Yom^rEx86NVXI8A6qjl zsv_GNCi2?4Am$yO^HvRfv~@werj1H@OCjG8DUI{PktvPOLPliWkUHSju;!+G8H#L&+h1iMkNn*>`dSfk*%o!8OLLawYRY+UwRgxsyfu+85_4C8V; zF}yzCK@8>Y6zodTKJOwqyoTRR4DH<`*f!K&Q0`NTC$f8m&ie$rU$6%Rdr+{41bbMp zM}l_d*PKVP9L~kO4f4Bh9u+ot?*hiqzlrMeE@Xa>gZB?^pYL&=O!m8(a6X<$>7rV5 z`tg*;x%y+|FSgsV%y6FUcI_Sq$F`?aUYs{SMS9J=%i2B9x~`jc&$2FNdmW8M^$AdF->VqaOSGrN%V_s57@y~aKXxyJ_YKeIyf~kCy|yLJyBk}AHn?qujIUu|8_ruZS3`w7yW-Nr7;$- zW?k-oF;&ZT{Dxrv6YNdFmI(HiU~dcdj$rQ!_MTwx3$~OPUiUvBW}z(0@%)e&&M!YA zW+F#rxsL_=M6geZVZDDw4DEd`*cXB=3z(Pv%Yb?Lz9NS0cR4Zi^J~Gr5o`r9d_MH8 zVBZl#U%nUY2Vywq`H>jvtXyKLwZ~(R*-uFh+t1IW2W_k*hBoFE9lm@XhwK+2w@R>I ziQ(@Cej|o{t|o@%`<+;kG-uQC2QfS6ar#rR=ziAa{t|M33-*s-p9l9!Zf{mk^E$4F zWmyid4KDX@vWIc1u;1ExFuoOuSx(neNwCVqP-hilc&({Q4EuGp%%U>cx6|;MY<1Fu z<82LMSf-kST}|_C^sN@j;rOr~F|@b7VD@dZb@k<2Q{AKe+CtC8q75|;?q|Gyn+aBj z>|y`QkgF?LJ;CY=)lAvZ$AI{Jo?*U_dz=e5NI zuzY;j$DGN>W+Y$ayt+!HKqC3k(izK{lN=OM~ZD90Y+1BjuY2MTsjX3>2A;KYWT81{qa!24ijugQD=jM++e|m5W_efP7KS?QH&cyh1@W~h6^@Au#titA=oIvqVG_6 z{YBrQaCW|!CyW;M#t0Vek92#{yFi>BDRhn%Y#cFc6Mu=Z=n&C98;H5bc(RA>Y>?3T zr?{>iC3H>@Y@%RC3pOdUXxx~Q{KRJ;lZBosf=v~yMZRy3=BIr{U!5j&9zzVva4a#b zvo^w)=|b)}!J>Vd-p=_uEyd;897@OWLgx&@PU~;d@%3@0keglH@WyhplpHbab90FK z8jyK-0x>)1aXC@2lLR|iuv3VkpQjQtUwm#pml(FCdBm^`rx7cv#WJL0J~2DzWj~PG zJoIOQbSxl-I=AU)IpTUhljJO?`?)Z)Xxum}F^tpMqz8RFCo^1Ya~pT0 zF%+NMouBfe?0G3&gv%UD$GKToP&Fk@3ro6?aIp@71rP0>qDUJ79mu212cp9zW zyWMF0-kGo8M{%|FdyMN<*?yGQ_8{Y&{A%H=t^>u9P*C*ewjMt?!wuNi6?@@egJ-(Y6&c_WRudUBx8@w^+t?a_-ulsN7 z^O$#W%8T*5N!YRVdX&96=Z#{)?+Ut&?A)5t*j8>4y7*l`cc#4kp=@`gH1>nrv#zN8 zwiloTWbaOS@%ipuLe};MV13@3@}eL22wB@JfU@_eyf}{Bmu0;@wgn#_5bQz09un+f z!5$IpQNbP)Y*jrgQ2Dh5%lEjDdqS`$1uGNmDZ!o=>>0tH%`EE6-=&H_0_yWQ(t~~K zd159uKFhu!*oy)4Io3-78>ycd1>@48vSf7}k^YYaeEPAJs%heHb8wmD`7;C>RoGxz@Md;1hbZi*t z^lQmR#73?WDgLcX%uae4>Irs+sQ3CpE_zRs+p8I@VLEFdbZ$%x`}-!uuzZ^eR$KV8 z8Oh-o-7vFgtlc2##D2Uv>FHL~$eMo|5ktQl3pQQ!)h0r&sbI|n+d{A{1=}jKC{9}^ zhS#L#LQf09S_;-mux$iuE!eh#ZI@Y8-YwEYhn7(HZG@h-g0&-7B(vW3#IQYe5Nvy5 z_*|wVG1OTq`p*tR?jg~B2a`^W*^WZzPJ(q3Y-hnb3)V%jT?C7MJLcoxu0n3NfR*KX z-<=qaReKP_`rA{m`ZOm-xvh!@Qhtp>)|GVP^RsTmP;NBMg;1`0Am?S@ix`%@2Qjqw zl4ye^LawJ^y#(tmSRcXmCWhCRzL`biO~1sj@AM}<77LHzKEyEg`x3)3dp}|r-vNS^ z3Ks1d^Zv0v$zfd|Kn#89LH7eFcc73vNU(ziJA@e88%PZOJXEm5h~e|^LBw!wG?*CL z8$t}_4kw0X7)lIv4ijv+U?T(@DcBLk!~A$GIwFwsHa{w0-p&S5d$IDl+-M;;MzAA^ z;Wc(FG3>YFh+(XE7i09s{VZpBTu|qDq4R3eiS>SzkeeVFze5kpFp=cYUh5s!E(6LP zO>)R45yR(WlZly}*X@*md7GF@tkeYaF^!m=cDd+1dCvONb2Zd?jL>;3F&wX_6GLB) zQ;eSDq1^E#S9Ceckd7I|?3|ZxreH@>AH()JOUOn0vfN%nnnU1tHe2YNBiJtW{a|)D&!W5ael6l zn-?(e@23UK>uA1UrweulF^tOsVi=c+;<=XXXR)!olKsht`^o!F&dfHdddBqE^V#7d=-4A6fIU!BtFzyFS_QUdoJU0J~L zxq@YM+tGd;XTBdNYHPM12kZQ%WCLS)V@l(7XQ}SQQ zD^ysV_11uMULM=OgR*y}yy)|tDP26C$g+2&G?wA^lurNsciff|Sf*Q38smOTkvVRA z{Qi!6QeJGgcW2pXOyK=O_lmsyt+D$=UfUmp<-K3zwf#Ss_ko-@>MORtr`dl|t5D(L zlo#vap_ImYcrfdV+MVt9L0ylNuE$au`_iK+T|DN__8-Z*qkd@nfw12_lkyhBc9wlQ zrHjW3IsH^h&Y(&@kdjwe$B+s6|HEboVH|I1mS!n1-sC)o3Xy&%|&g1scz%YwZk z*sFrQCfMtOy&>5D1bb7kC4#*r*xQ1=BiOrwy(ifFf-M#717g+$eSG|o*kov*9}#N_ z>|Cx-p~3u3O%)Gi~2<@=JD<#&5u3AS9YuLX^mX%ys#^KM_MeuOgjTe?JSk=-rMU z>*)QB&Q=C<_Wzi$=kL0xXwK)u<$e)zUltdtxQ>?gGet3bk>Y}Lm{r1Fce?*G3-0I6 z!Si!x(Ov>)zY2TR1^Z2~52$?T=V~GMJFy~Z!s+;fn4R-7{3+O9g8ePnKg33w!F>En z%uc&6HG{PnXVr>|j@!hViAT)%T^eYQ*Y-zbAMH=YVz4RKiE|}psIwBuA+z=Lb+zlh zqJ35tI;#j)Rj_J;RVP*?&GmHDAU1i;>Ehp-#O$QkUoFAb6Ks9K?kXy!ypGUbZ6Q}j zuW84420R?+NM)xdy~=zP~XsjP)jhohQcm zO@&-%F-~kImt}L#3n--b|q%W+uCjc^O)@}*dBuIDOj{t>Grw` zxo(0TpU0+3P-fTJJ;_am^Lq(BBgJ*Lhmbo~j8z-;HRJKPSy!+UVQ*lt#^dFU*6*D4 z6gqneHcrH)w~%Wh%GZ|qoJGm~jMnI!Mehr6)<^iVw_tq*>nB*W?&rSDELnHG^S7e# z+TCB++efe$$lhc#pO1ZooccZeenQUI%%Xd2TQjpZAC2#}R)x*AH64`l;`9CkvuqX6=XGgoP58WgV9MJK@*R@Wc+Y>Z&}C~) zSpGpNFP7o3l*V_D9-4JU@vyZc)HNjK#renJEE|oBwzgE@T#mzq9b0R{@(xRR(e6-T z$JUFmJi~L|sNF>GTlYG!^%`s^BZMAXpIKM?v-KGC{fMLo+tf&**Wafe)rGCcpx)6* z56)voWxbKCt-+as2=(E*JE>Dk8ew8r{jp>`p0;|juLEwV5b#MlwbSM z-h9C(3Y|v_cDm3xNytqWY>Hr01)C<=F@hZ{*mS{;BWBg^edl<=W(2G(x0#s%8=KiI zVpH;Y`!k!EopZV2!LvkXbA;T=;QrJ7JVD5vDA-BFunZ>?n`DK^$0@|@w3p#j!Io3m zk8*|S?_41_Pq62L=XdVQX+mzkV5bwC27aDF4CAtZn2EVBX9~7Zu(Jd^Td;EkJ6Evt z1Up}_3k17RutkDhB-q7*T|x}60hbcPI=W1-%LThauqy?-O0cT~=KcMefO$K+R9WM8hkel9f-F5V`kb6b2R|R`bu-659 zL$J$4JNuuIdsDC_g1seJwCBa^;cX!|RFvTzAvdAby6f*z>;X%Eza@q-`%bX$1^Yp;9|ikK zu=Rp<0x$c|LT;sCUx@nqMaZoZ>;Tcneid@RDHizY^ZC_6?svid4s?1v{}6J23ig*^ ze+%}HVE+oXX2W_szEmtGjxnnwShQ!!<8r;2ztjrWY+bIhuvbN}s)F(Ni?QEU6LQrB zt07oT!J_XFd)aFVx%C8FKVaVGYX{8j)e&q1!8Rn;62{1lh+(Yj3RX|BmchE9m*Hl? z>I1lw4!O$6Iiu+0Q(DA?wL*}F*Abzjv;$Tb$ML_FtfBIKG17QH9L>o3|n=B$~} zxrJa`3U<3_FLw~b^|Gym&aDM&E?5h}S_;-mux$iuE!eh#Z6{b8Vv}JkXiE&of_8#6 z71y!$LhgLgcB5}sc{}SMbZ#$LN5OUwY)8R%60DP8I}6rXur7k_BG|5i?Izf79oHVS zu)XXqk%;bvm{{dXHUU;3AVnNXZ03xeFWQEu)c!z z6YL$)H~W)ZR~SR~A%0#Q3U-)aQ|b9N&J_m^;@Rn0LhfwA&JpZf!Oj!xe8Kh-?dJj^H$=3p3x(WBA-71# zT_jla-X(7%usI3Y`}VKTj5Nmk7B_1-neJ%LThauqy?-O0cU1yGF2U z1-nkL>jk?(up2Xr#s+&|Qcsu@-6Zt*_qTBWthW){7rr;>xjE^}kQC5$vSuYvSRt!ZIupWq4NTd`__E z1$#lT7X^EXn3d?rTo*3~tSqxvh~aNcUL}V2BCioMb*}UEfVnSk2=+h0-V|(!U~dUF zs~F7JUaxdN-zGWy&fPo2FfQ*Bv$(jm6XQvR!-^S(L>f;DW$F6quAN~0M4&#kZkoV>#;LU?x)=^tgAiQ{eZD3 zR`xqi^nbl%qmO0JN3E2`c2qO_5b5IIdDcmJF&4E`8solx))m#a-Ip|TC(_>ya^9+N z&dbj46s?!?7E_6ww{A+~J;6pPjdSP?v+gLa{+rUsuHD06o7^PnDFOdBPU+&Zm(Mjw zX>1Giv)!mI*l$o#_huyb9L9iVK+exra zg6%9=XTcU1gF^Xr8tbo%klRJDT?N}su-%DS%PO;f(y<3I>}Putv;1T1pLBE;tearn z1=~xo9)kT$Ys%)I`&kmmdD(krhGU`2-IO|SOH-7OUZiKZgY>sIF+;AWb8tWCtPjay z?Dr;yI-3?3+VZxAtgn#kCs=>MMi+x~Y;PYSx36IP1*ex5`Ouy z&)c61GK<>Qg^8hWi%5^jdkk06ToU`)MM5t6_Pgu6SjfFfF+^W%&W&vuS@b=5xA&&d zd5Q4lQo$}GR=8#krQ>p980#yD70ESX6|N+Pa#s<<=RH>wGj(3}YXrMiuswshf$O|Z z$Xzek4T4poxhs~TmAFURB)89~o!^-3wal@#KQ{^cBgEWmv5@26|DxYja+xCE=A}N5 z^?Nhfz%k$!Vp#XL5<~k-g8t&yjN63V?Sk>U^Ua>?yhF(GZ;Y|*cM3UvZxzbjCFJfF z?DyX0b@_N}DR0~N2)TO&yHBwD1$!W~sEvG`#w~ni_@K~JF{o=V(+9b2M>ZZxI&n;Y zSlFl|`e+p~UxUe^xEE(#C9STBGf4S@ZF1_9f~A zHB#B}I`JI&jbqO9#7w={*$aZbDA-Gay)4)(g1su(YXS2y{Plo&`Q8Ya%l$84F88Kj zO9XpMu(vad>hqn%us+`vdfpT4eZgAgwielMmHIA@sY`{P4+Lv1#()on+(OZ}J`!>t z3%0fBlXXO&{6y&dRItwk`G}Dm&EXz{1q`=mt9T_efgRg zmf;&>_t zP0bXi<0oQv&g*kGx^G3fpM~5?!F~~Jm0-VS7S+pdi51fab3Prbg`VF9>qlcn5$-iA zEBqnk{uJyl!Tt_dS+@62z`S1m6>QDs?D95Z>Rhg3v2ir_plr`>A*h(>1YboSf3AT-3 ztp!uRN#0h-Z6{cpfR*LCXd5tZ%k79^y|*W3>RhfvW>NcTnUkditmG7E{oiw?yQflkO$)`F78ow_=t>y(ftK&R#<9eDTb)calS2`v^Vr=sJSF z?JeY@?`wIiqVH=t8<66N_ho1=`u)7iohssV9NinB&c31y(eLivUi8~`XGc;S#&+6I z*y}IYK7##8eH-oVE96>;=eGL^xdDQe3KsoF-D5UDJU@@VE#`9jXM5g%v2N|R?RXz} zfbeYw)fqm&I#AfK-^k;baS++Sx;U6v(Y4IsbR0s=&UqQ5-`YDHD16(ybnW@#p~6Pr zly_bIZnjtpo!onEIULsx6ZQrPHdwGBf*mf{iR7nw?JBRC4BZfLx=keFa zbZjHTNeW3~IYPCgb^FhM zu#|Pq%XzE1k@VN?*l+&O&KW5$#`^S>E(N*yDUH|2)3WZU4E7tpVwlVFXOgZ3S=P&d zb#dsX6)F^uWwY!;VaI;+R{}nso$}&b`m8ML^FPsV*Wci zx7W?G=HnbHGd^cOFENbuxx#PzJ(}6C;)>H>FO&VA4BOd-Ne|ZP1u2bo&(C(e9!#a{ zvfr0s-bFcYWY^x`fp#t-I~S)kww;TzE{{`ND9fdST_)Iz$p@?#)Oop(yF#!l1-nYH zs|CA8uxkaoj+m8pwEdHg>jk?(up0%tNwCF&-7MHGf}IjPfA{j;D&%ex>~><TEgV z`cwOx9~Hj*PV39)OSBi+W4$D}H*oe?wimU@R;j&UoE|6Nu)Mocc~R~Ol0(0r6fD{g z?0)n2f6$jQ(uwn*r--31(fe-P-p;|AiL>X62Oe=t-+?-w?CYF1JFkO8GkIu|ToLP9kJqZUiR&KVrb(B!M+sN7u$n?eW_5B?focp4i26(dYpbDIqVZZ6T>nz z?!C4z$VLWhK5lQN(AiP2Uj$nv*sp?(6LI-X$W5m?B$j=(ko%n&-f#aw40R3{^SwWX z++TwIO$^)oKg7_Ne+7%)Q|9%)56PjuH7{c-AD7Wy#o|RxG2?eRWB;opRV)r8LK#IWDiAckeADOfGR))Q=fV(3e4V(3d9!8Ra< zIyWSSIyVw5dSA20x-Q9~y?VsZUVUOH*MJzxMeoXUUp5wU{(QQ!Da(h;**ov>y1Hr7 zgTGtfMCj$`*QVO_+B0gq?nE~28FVkh`DmDI;JCI~wi#V(?U^*nZZ7;=oPCJ&*z;$c z7c>%jqIb%9IsKh-Q914TF~+t@vQg@VPJbH<7gOD=t7+B~`C#va!`N<- z^5UGbS(c6L*t6>O=5woiEw$>?NqjKHuELdm31{Do*?Rkpp>_T$Gvn~6x z3$bCqb|q%Wb?zqE?t(@4@vd_ZA-AVsT?OkVSoeUHI_fFp zdI>f)cs}pG+(Gw9Eun6E2Rgm%eTZSZ-kTV) z(0PAxag5~#_FT8leT2?^iQ)BkKVle{0fLnZwm&hfhpJuHE(5XyNDh6uQ?!=@h1{Ux z0vN|LTJ!VzJBW1RT;pJ3=x6l)8rOMV1KVcsmPIB1hM-VepUWN@w zC$f=3?)Tun*>#>S`oj@IXGIa~Q9^DsG3-ZUh++AT6s)%JWvq}pK*aM~8q2Yc#*xmV ziI}tL7*EX3d914k&*Gg$?@MsDntVZ>(fcP{?kM3)eG#(>LXOw5FfJ2?Tw}6_&-afO za<9=d!tPf1d`!x6mHb5d>t&yk82UCju;JJAsl*DCS?4rjcG_(mLk#U5OAO`wy#vua zVLHj7&f|!o&dsQg<9K>J$#pZqe9Rzbr@ee{im`U4kc)n!?>e_AZivgTzqnp9OX!>} zShT+A_I!QOV~z2gBXl+>wMdlviE%lB2Fi_~zKZO7F>g3m*gH?K z^98#=unUP{dtF2f%izEB^sxciMMCZ;QTF3SJfq(Qdc8;Aly&CcagF-N-sxc^miJ=v z8?R-T5X16bDp>S8G54)L^$%>5mkFJh3wD8McZ29z0G8nj(%I5rKCUEYr@ahU33jz$ z*9aE<-plRvDK+cm*A}e5YlY4^;`%X@+GLS5GwHaFbmEwDJu%Dge(osdoi_-%8wI;b zu*HJCA+BW?iZ*d0wI!Ty+$`)>rC4Jf-6G_8JqmxjcdL-QO|aVqJBrqWvHoiIv{ZRJ zLw03;ep1;^=ELjV*1xd*+?nhb-8|>JL-^zCU(r})-&w`)pWc=AL_XMh7}nK2DQ{7g z*@wHcEbgm!|M@o>_ouvA5BH@smgnBAEAqqE#BfY|Fy+ND^#LJkYh;-B;glEecOS~K z(e=mI&?4r)!H8mPYh-Bi5n;oBlM(5$HL`BrFw);gg&zO?Mpek~bzy5@INm&&^c0uK zocD>8#`)djDP5S#{IMbtw_#f|TUYz`HL|EqZH=rL+OmzOk`0VmS+*IqEn7pww(~5> zK9kZ|$4_Tnksr3EhHd=$lo$PaF3U!3z}MU&z5HF+7n7c1I-1M-0_km=&)c7u1bbPq zR|I=iu-628U9dL<`>Y69d7HuV{ZGifDcBOh-V*F>!QK(mLVOV17wikcKHAPSm$w&7xz6Ys z>TDV5oMeLe_>!2NcDb(vTQ1nwf_)>{3cu&DL-$L%;;Qqwr{t~~D-8Yk)qof`n_+{0rP&`BVgV?N{FGKJ&Bp0uCrIb zTxV}$*yj5X!?N!!SYKlJjH@3p)Y+dH%I!l8<@P0paoLX;$_*ffWiKU$I`*iEq z=;t)S_7C;|y1nTAmClY4I`408*5f{XtdP5y)<&@mLq%JfE_5D84BN!<#L(Uh!Db3J zOR(9*u;0!hhQ6F2*olIjL=4-@$;5CjehM-3#pevCW){tj=Ou>mog3)!_H&wG^94Iy zurmZ(AlR9L%@pl!A<5x7+gZf0J{MH8)Y`wh^tjmXK=Jx zKNqJo?sdB;r7=#6QX0#1VM?!S4D#QcMsc$Bb@b!1qzB`6snF|dsVrBeG~QobDSWbZ^L52DTK9GzeBC>WiLHBM zY_Ada{e6jr8|F$n+=i`N=$+yHl{c z1iM?Xdjz{zu=@nNU$6%Rdr+{41bbMpM+AFRu*Zm5jg7Q_((yR4acfQ&|2{#?PP*Ki zG;hLd@RLHWj2MmuPZ2|V^XdKr+0!I9%3wa8A!et&4969ns(jAq=Pl2(Lhd>8{Yxt^Cv4(sz3VrI(aUKQ+pitk!- z!FrPGd`;+lU9dL<`yVkZ-ce?KOMars2BPX+r-u+IhiLa>>l zj+O;-Wx0%B2FzpjRlru7i;KR>*xv4Ey`{#IQ~L zAXxPKbuasmLT1Y#zXbbvThmfL7r^=T-$L#m!TuF&%_i`**ZaIJRZJViTbeuh zs8nniFZaz(d)~@|wa#NoG;gSqXiVkozW`ULjRnk0vF z#9G8qZardX?=otm$X1HExu zN4t{rw|3TD5ze{pTZF&SZ~R@ZRf#DpZ)1g;tbZe6uP!m{`}K&KI=5F}um*x{EZ8Q* zuzhbz%zW`NaWi5?C@e!d8WO{M;mwIzelK66%%Z+`dvXl*G!}a5s_QM?GvL_TgmmIu zu_-b1`xnuEnhCir0_O4EQn0N8<}ut_u;zkA`-EI)&EjE4c{|6l*CK}PpoOqkmDZI| z=i@~qjP11)I$IIL_OcBzwAX;jfUGsiVPD;r7|N|r_xkv(ayuc{MzDR00dMWNczbC} zayb6CBZlLDdtz9I4uWk@40U!ShB|i;Y)8S~5^Z89A=il*&WUy=hJJPyY`D1Z=t6QB zm&v=Dg|+L?`&jf&ZD;q0Ywa#%56iHtV7m#nJF%9gG#`5qv(xV9o`Q81tearn1Lkdh zuYkEP{ff>X$GV4*D-o=xV7-W8z4s=DG3z7P-hxHnPxrF-6>|NE;W*Ks82Y&nF}!Z> zOAO@>7V~-jE;-8WCv*-FtW>c5i51=0jHcrNVi=bLiCKP+HSecEb`Z(+Fqn^niP>q_ z8NDyu8NV+aV1oo3ObqL22r=~YaKVNOHjEhd+u_8} z-Uz`)3U-8Gqln?0VKgy}%NSxfM>&!hmSL=5+lhO>aYD|%v%aotTy=2`A1`zsCD;VP zHWuyrYPxsBxJ(o}`P<$oceIeR?_RGfE|W+O>vl3Rv=@CZ+Q*40B!}&KDlybKjTpAU zV~C;Lv4Tw!*MR9l?l@w2Z*hEP(Y@k~#BlCElk}k97f?IsW-uSKgdD$5q|}Ad-`PTL zj$jk>eQQ*vhNo+Rwq_p8^{ zU!v~-xG&MW{+&hpe0g2Wzxn+s=C$jL_SZQ(Mc6x4upy!z<_fuaf<=4N+)v({ zhUGg==$s!gpSPbLFt>MxU<-(0TRM{%`f1;)T~}LLDCA~TTO$|sN86`WEDc+hbexs- zM0ssr7LMs>3!C3xurg(<-`8)_YOPi`Sw>SvD%Ky$1!K>0T`I+MX`FR$U@|viG3i z{mW%3FOJ`rX4&XjvF#6AS6lM^U{MThKUndYDaYXovU$0%>F-I2{E7B^dA<3bu1L<` z_aDXG_HbeSUYUK1%4K`HFqT(Sxvt8xksaHkh4p)luw#3&Fn-siyy)Y#S=P&nd#bJ% z>;}PZ6znF!77KQ>V7CZ%t6;YYcDrDA2zIAncL{d4VD|`iuVD8HcE4Z`2=<_04-vC2 zTbBFn!vXXD_K0AQ3f8IU1Z&U9Mr3=B1#-?F7wieaUZDBfG1->=c~Z!g33f(sjdede z(;lj!VDBlR^J!vu-FSwWsdRhK3ie#UCS-q}C+7MLzd+2+Ie(E@F-5Zs>3E5lopU+f zw=~IMK3*1buLxEW%$wa_@62#4cJ^x0iSd0+*mzyAHx#3GXGL}U{}Xa=5}WElN`IFS zL!HA#EZ!oy@u2fS!QK_@_TawD%l=-Li|Xb5#PIq1Qeoo*!9Eo1Bf&lv>=VI0 z73?#?J|~84@(W^EFVn^NvrNc+Net)LUlBul%LVH!#`CX*+&6+>t5u1kYE! z9{v?_m5M_UNr8+8xm~3M`(t|c$2%ed`{Tg}iPh?~Lq!X`?wS|p3 z#IPJ25JTUpi#cjzYBT7|hC=5?f|ZGJdA*`h#{HW2Gj+{zw?B1-y?VqjF7=6FTpAE7 zhOumOV`8Xt6Jn_IZ845)D&#g3jK9f(ZL*<|+eq}!&4pa|++O`T7nbiJDj&w>PH|nm zmF~|`XCvWDy*V+o*FvyY#Te6) z}`td&M=lwXoNg#y8ZtZ6H@> z8Pc&`X5Q}XoX70JR0bR;+XQ+%Z(Cy6UfU7FzSN#rsX3634#e!V+uNQP>U<;UpI+XM zLT(3QXm3Zsb|QxNlbwj6&YcD8ELa!8?Aa5}Cp}(vb{8SHYrs5ay9u^?z})5%F&6G2 z{AK>n_+{g7pxrM6l098F~u264A$d3Ax^a^$~0{8i(;Zv$v4zE7<+w zUcHp&Z#Wk86FU0~cBB~F_YrdY62tLjKVsImyv;uzjNjgm1_-%Q!S)wy1oe0H-HdGsPkaK4k3nfn1RHwZVx4f_I9E+i0m*SH%PEmK^=MhH5K=(gN4rXh0fko zhN25tLDDfq=saAo4aL33ufiAJqh&sL*@p^ye~NzFtEkNQUT3-(v&QG=Xwlrm_9g9Q z&gEm6@NKrRF^8UynL01;aG`UAU?T-PLa?1h+v0B)VP733bdDA*+H>P&w>>w-WhlpY zjL>6yZg5UvDc9qBY%1Hyd^oc`HaG^1O*V?#s+{l0Y_p0BroS$0dv0(Jc~r_rY zH&ggzdxNlD%psp<3t8JcgmzC%d9fXwAY}RdP`9NRlZc=($#~Nx?kBZJb&RAmy)aU7X!Sa$e6SznB>6yjieY1iO`(`(e6oBZl^F zCx+LxJBXp&odGMe4C%NlVD8V|#Kx@=DgM2On4NUFdj-2su=@oYE9&n7AvcZI&+(e} zppbjBxagL*T^x%Z5^@g<_K0AQ3ig;_j|=vMU>{TY@Ou5EkgHl^HC-MTtar61r%dR) ziN*=kxk!vJy@~-l_T{T!jojPkQ^J?01><)j?PZSU;~61$EtLW1r_Tzx=ZN7v@Ofh9 zhL`UJ!CoYW*SwdAq0W~Dn@;tB2{ts?Bje@!Rh0euJQhW5`O9Pv$KJ1m z{pEt$dq~&SChL=Lct8HN(0QpCo0ib^!>oHOE+%`({-O4R?XE2Mt0+$2BtMI-X{Ib4 zE68uu^DQyU@4kJPS!Cn;#9CVBeEdLq(8iC%OvG(82-f|aou1olWZ(Ae;kfpbu<mTZyv&CdwD>QFNV)MZa53Ix*J26T|W?5%G-POYFY9CFK4P_IglUx|=ik z_><(Y41W-WDxXN5yeODy+CF)5GP z0E!d(a&IwNiRCJiJ&a)`VyLrnz{)H`I;sd(Rj_IS^E#_8SPj8y3RX+7)uJD-M{;-_ zT%Q=;N7W`~vGVf$A?!^OWO>L=zV>z^Kap2{XnO;r45BI z8xiYf0nJBUVs_ensYk543#GsHiJ@Er!8R7GOpL>u2)QOS$H9A!O-T;>?PkQV3=IX_ zT(CxhO%r2JVX|i;h5Y)_|j6aRsr*I zX`6ug_|jUiZ3Wv-ur`9VC5G*@9WjjOH)1}~p5$;patC6l^D}C@r3Uk{y^tG7eFew5 zjwFZW+ku#w@_N{j7|QKLthk-Zes&^;I(HVVb7qyY9+YdF+AiL+bP;+Qi7{*!A-Ahw zy9wVqr!fWRyt|Vgd``3nF)ZU9qJQowv;1RE^akbwEM z`|yDIHRBP{SB46?VS)`8EP5xf`!YhvjTGz%!A1!-TCg#K9Z3wY-D8Q_Sm5=%RZeL${j}x z<&GC@hF~)V`p`yFdruiou?4PxSX1q&%27pMOo*(#EM6zIX##3 z*lCZ`r!;;SQ;{rpnvk0>*y)0uA=m=J&J=8+U}q7-7@ke6doH*AIfs~?^O&70*bfv> zlv^p-d0D5|C(50l81DDDK-j-futkDhB-q7*T|x}|%caC{EWC^umY08zWz4UEmkYTo z1RFxnO!40GN+DNGJcqqX$Tbn;YYozgF}qsmyhgBV1-nkLt%RT7i?M#NxVBs`?A;*P zjl>EQR{80;i5T|d#l$SX_tl#PyG5{D1uLcT9PRBd*lnZ}%WylfQd63bJBZn7FW;TS zQ0HC5P%ipjfa|=Q2q|pOM^84E=mSum=SjP4}v(^PwykjkOOa zhGXp`BCqX<#b*&?#r~v6$ycl=C3gYMjj$Y#kxpcf3-$yt9G9OYhW93A#L(VTnME;u zIx)2IjL`F}V9yEmykIW~_M%`f5yLurnHa|6mCT}8987KIRnmhtUK4uk+Ya5$Xg*#i zJ$Bmr-W!6o6Z6*pksSI~bx*SpuMNE^xVp4Yu6oY&E(f_)~~C~?i~ zS4@uLcJ?{xEQ)ImrQ-`?cFz4=CfE<6ZfnzhBkGKP7wz^o5_0_g53Khug`cO2x}BfP z5RE?vqTJAd zuhH+-U8jE=q@p>Pjv`EJ{8vKLMIa2hwAB9{C(I!RK-BeMF@`P_awF+}9_#wA>`T<{*GqMc z_e*~ZzfTflz&}FnU%@IIX04>WzoOq2iD7#hL3>J3t`f;%8>uW<6=GlAK z&+Wvxwt>*OA+e&%#)E9296gyYKQLQf-NDAzc%$i^M1-C_S}LVCJeLGsa*n4R|a zyP%q-%EusN69x14mg&$FA6uSHkLI=0L#iU+?p zXRDky@`1li)jZ|x?vAFvTMIk$6qLi)C#oWbNHJxVG0ikcVzH^|pMr7er#~!3>uaw4kbWdq4L$|CuDucbt$6^=No4u0-$;m%n#P7nkas?v>J5=AJ2yewC#3y4tOMHwDo7?@p`pSN*9;wtfPNQo1Y%Pw)wpMDHUwn=9VLt+h52XAlQL|9VFPn zf*m5*K*0_b>@dLw2{u@;A%Yz)*igZS2{v4?5rT~*X05>6<`II85^OZFmZmfxV~E-5 zVVNB%*eAhr6_?|86QRzrfzGmQZ(P7!=Xk-cqPhCmtk3=&CFFi629)wXX8GM-^e!T2 z{JtTtACsRT>`fHxXks`{Od^JUP7avIa|$sWU#Aj7oznz6MzCWAr?j5JP(>5<`0@2{xmcWW{~wWRmL+BMmEdIm8p`|{v9qu2WaA$O)=L#YgC?_3%)ku4ND`5h`KH=&iK%43G?ReJV} z>?~ogl_#OpGW0P9ND`!p{qY zy$c1a9z2`#wp)`Jjs=T^&OY?a7|XC8=|r|tT+{CqV^*EuS(%q1`pv$x{=(0TMEO1v z?K=84liT}NwB3uz9^ThoLJVVeDKVTMUPjF1%Iu$XTplp@^LOFr2Ej9Tmy3Rz?(7O- zuYu56NwkS8h0ZqZtd!+_y14w?sdQW=X`F7(;Fo zI&T;34#Dme>~Jxc+PZjP5x0rENGHbgZekeE=vzb{YySQZUIXqCI(vye&EEz>KkpSf z`8SqWhWmsZ|K<^&Io>bix=@?IvGoBVw~Oev`--@1E$+)-qCM9bvj>GQ4+-{LM~mRv zZK=#orQ_j1&e0OV6PIxxV$D<^zC!6+t-C$v*N)`c^^YR-w<;D6YNc5 zsB;N1EZJO;%J(5HAJ~6bnRIm>OJDJ)UK7;&F$VK~1z212* zDPF5S5;_;sd;-h`_p4yP33i?6$E!)Mn>m_~--%)Q{vd{J@=s!>&g-_SxL^89$o);M=n~n_e~9(W z=k3qG#O$2=QjZwkFI5<1Im@rV*mf%tL!Ffbt1MU*!FX>{@hHrUrlV?>tF$Ir{Ojej z{YO})YQl!?N3!xpdh9!0_$;8h&|~|LOiyL|lMnZY_Z!trdeFZbqAa%G2Eba(=IdqMQ5kGM62@Y~l(#g?*`E!vY(+Td ze%XE`i%sz-b7#E&s6o<$v8hjX>ZNqit+SrGDP2sPa(bhbURQnEek1gGqqy7t zqIJc>_6K1snq)ndOe7y(9@{U3v1pp}7V0T4f4ghzl()Ej%K5fR=@K`P{%)Dl*siun zY4odEN?RGC^7;Ot$Y=hR*Ed_P(G$hk_VZwyYmsc?eR%WiLsZwcp9f{PNqMoIv`T3# zPs^+;D$h1G*Dh;YOXj0B$syZTu*JFg6$&Mu7d4G%$iDBZZEqNo1D+vpFN1#Iq&Ox2F!JK4cJfn7yR zXk$6o?%yugUC8YvShHZRA!)AYfgy%~8cgwcJlE`+Hr01)C<=E3^iJ{q`6kcdTI3iQ)67+LLJSiafBCYdw&m_y7?d;2^=FkYuZxf6xlNrIg$*eQaYD%f1X<_UJ1VDkk# zU9d9*+k@IC{zhkkkULYbvBC3vuiKubR%7vawovHg=g)ZWc$SbmTd;EkJC|5jb2}gB z5wp`i*FQgCUiJ$Fn@Zyu_M;1h+~vhXfHA!EO`mcEK(U_SJgK?jSiF8}1~A<*OdN z@7Hy1pUdv;73JGV6R)$VL0{UFTy$ZZoonzC0=F z?{T4%zk7%}qy7Kx=MzHbVsUMGlH{--mk~o>o)T;mYM&?!Cn&VWx-w{hIRca zF^tP{F&}tM$kn6giKz2+As2ne&SM>YtInB!k1iUQ?YnX~7QP|u|Bo1s=Wi0jv2Y17 zlzWRyka3 zqrOLS`1|ztiD4ZrC5AdbAcoh#4~e19kBDIzrqFm^Oj&bWJ{EFKD4wXZsu->%wj@jQ6L!Cbe_9HR$^Cx1c z^XGuMpDP3AZDJzrCqkXSkX+YXZu_%}7|KQ8K6HCK(O86XzY3kvZ=77`Z$gg0H;8lm z)j}@%-IUwgv>3GFHs6~1EynD3vWH{eAH*;&e-guM*I&enA&`STB5Ne%?xFIbTw@{EiRu>R zqVFEMpG`<7_LZi@aGuwU7`B%!h@mfA3bumAPn6qA$ZaiHbHQ2&){+>u&sN0H&t3Ah zGP=%eljKlOYtn;ld0WA@BZkir+Ym$FI*T=;wnDBQF)TxSV)(qg12Obv30@?MiwrtIKxJEULG?62tz~gY@8iX$dj3 z(UTa?U3z5}*r9@p}l>C+`fYCCzyZ#t)dyqhnJzZC||px!s0o}0Aa6Gu!BV1?kCoK_Rl)K z-C_M5m>9OL1IRx1iGzq)I6M~3XiP$#2a_D$*BwF(bq*BlNO6xD{f^lE`&u#HS2hJH>KY>Hr01)G*x)R&G)41GIR=ow7?75zR$jQ2kl zlbd+FuP*LW2GJZF{hTiRJdPONXC6-s%P@l&_MMr6%_4?6XA3q*uoH-(FDDX1do#t@ zf0B@!Antj`ifil1q!a7t6k_P-5V418m}qCG3Z0whu_$Wi_N_9UW6vc$==YsPu*=&k z#(o~j;WhF!VkkGC7+$|lCx&qvE@CmMFkbE_%5@QID`zNs)aOuV&wP!lXeaaGW8ng! zvkUcqjP;pWE~>NMY2JkQ2n&S`{{~v5XKv~f81J)44?bf&n^@5;a#_zIhIM~#W|4jS zUK+NA^GHvr2Q&RWpBUczT|mslJpRXvws4TBx0xh|b#tMxw@9!HsJ_w9A;Oo2bS*$% zE)w=;h_U1<+TVfg;9}B=^>7I>EW^pgLyWlJ`TIVicF-sFNvzLHg^kOI6}NWgL^>`f zX6L*eTp?KW4KJ4)FUFasR4-`nG;!Tf@Bg@x{KVK_m01+SiK+eIbJeR!5Bk=(P#V`0 zzheaDt`Rz;?|pd;qi=dSyH@DDju^(`dSdA34T9Y$*iC{hCN|j&=Hq5!cG}CgTJ-r_ zgxsxy-6q)Wg55z3+smEAunczzcDG>j#Qf+UlEeOgFR^Yp&i3a%Vk3avPYl;vAIPjS z8p=6a9G&o`p}`bd()HuNaz!Mb^j82a{jW|55r z>G=-E?g`R^zRsj)F%zKtPm)}5smi=e*tc&&;eF^+qzC5-PZLAmo)PR>!EO=P)aQiU z^Mbt~n0@C8*YvRrjl`ai+l8O}&KK0_-=m7!w7v5M>;9$WH_o+Q6y@-Dz*ID2`S3Pv z?|`u~CiVXv&g`8pXyf%{1Me|kOX*@7oh4sQX&lpENonj8FQ>HGs`;O#^1j9Ika;)d z#qz$B(&)?EDUE9^Z>2Q0<0UDLzP_2#*mnPy(nS}^zPypr=3Di^H+x47`tniAi{sme zDP00`AEY!sn_HUF*#F;8Y4qj2ls4N{K!?`>zw_qPl((3|*Dj2Ha{x| ze%d>3FqYqt?AKW~DxM4!=s9H*6 zOsb|dwwWp^y{>*^?*KyE>!-XJv-PrURL}MfAe7x8<;AhLPD*3@s-1OJG1Y)&_hlQwS_f?E8kH4pXm0;gUIy$d zSEhc4%u*xmpLA>+=ybMSz=p38tI#H3<1%X-u+f>dCO_@8m%Uve=f1QL*il($hk#AW zYeN zcvz~~K>-^sA|N2WOOf7t2PvUQ??{(!0Ra_75Cj}W6tEWrq%R6KK*R!~D8+&u6_uiX zbF*{4eR-Tc?~i16esk};@6BYAok_C9Zi#Z%iHwc6UcxfQTbJa)Hco$DNq!xR{kwse zr#{JpdKwZ#AB}`Hj{3+Rjmv%gI3(O3a!1SU&VIkx`#shB9mf}2JT%&$(#IH!1>U;; z@t}X?jiP-CakhAuDQtB zBk{s6q8PZAu<++$00N1Y)5j!+6(JIY_O;W&jYHKjwC0}xv97=gnZ&k`x;``^0a_1Y41+30-wtCn)lUh(`7hzq6-9{`(^cIUfmpDlCF?1twsHHnG z>o|GN@Vd7sd43pctB0*G z$%$HSCx*Cw#DYY|SZqV$WPNfiux}DCf|3pW^(Q%D1Bl_Zoto=`Bn~+T5yP={Ffqgp z5jIrVFk+a`;lf4`!+eeuHcHrO@jP~aSKE+$8;vG8d)w+^W7}b5I-O_OcE)#*ocJ4Q zjIgo9f~dc7#IXIw3!5NpA~Ezgi5SLkCo$AAS=ba}=x?g9X~fXqbYjR^qk>Id#&$>j zwsZ!GL(ZAR0t*V;UBd1rhCc5hhW_p)79@I$#kN@ME&3Z<($?VpjTbhZ z7+#;&WX!Y=`9{?r~uOF^umCVnJf2W3jKu_z?Fb ziG!^mhMeogxf}qo)Y~%OFkonI#-Iga^ku1X%Y7fG4!{J80Mvdc)o3$ zlKH%>COOgH8e-`0M3c;OJ}cs$6V^_gU+PW_bv{pW!d@VTF&ra*|G{1qaofpb3~`0L zYb-WPuzOnDZ+!oJN%U7vJRhwQkFl3U&R2x(5U&ScC2{z>Yo;LBw?==d)1D8zQbJiNkzuByq4!!Zs5_ zEnUcS9PAVFI6{A0NKV*RV#v8)JXW@g?YDy%`r9V@+fEE|X~b}@atAStVV2l;mkTQ| zj^*3P-#O^-O;O8R!nU`Mg)RBGgt6}wab>;t&|4P7?GkZs3%gm2eLERD`g=#@JV5r( zz={o9Epcp}BOcT5lKxQ3d&1r)79=W(#Y&N}}ZK0gpS_X+z@*m|;0pw9gy4rBjF*vG;?A%^}wC5HYE5JTKSVV@C0&O^kI^90%F zalP?#5x0>XKaukb5qDVFm&7oJuY?^ThA|u!c8nPMJ5CJaJ1$<+wk5AW5!Y8dHy$M0 z0dc>TiG__HcfTfeVhrC1`&QU@#8As?;x+AhVyNYW$XT(%`QyX(`<}!h=MTbu6!w#_ zpNV1YzYs$$zY;^-Nn!iN^ZU!wLkx9(Ncuz0vm_2V&k;k;G_oD=zLI}L z+*dJc$;Xf&u`CwbMCKE<{7Z7e{v(E3VsT5J+#lky5DRSeu%*Q9HaRY^MPWNiwk>jI zjoTW@aj+o%xhge{%hb?d1Zh71l!mcD1gqf`6O|jeLoPjM0+f7C6 zHrX&Q`Qn*deqwlUOaWqnr3%~1j&_^eUto*E_M&+1_*HD%zT$krRdLJA`@5PL#(s^k zYl#K1^mZ*dzF=Oi6LAHJp_c23-JafEER|Rg+401E^4O8Hkcd0m&;ohf4I*y5xc*RB z#7!jE6oXiLyHUh_Dvqs1NF2sqlvof;Z-+$AVj}J)VIO(VslB4U?5?1TQuLcC5Z*Ldf0Z6^C&?=VzGC`Io0>Xd6b*umVnn%idbN)hplwn zZh2f8Vu7t5wt3n z+v;Jf7`NNxIuTb%SY={Cn8{j|xZU!cRfz?*de|nD<3!(didBo-4?O2B!~$DAY}3Rs zqgJhpTT-~^>kvwOCEehL+ zMCZwdxO#D0BiRC56t?!VxaopD6^Lfz7GsHECpKpOJ3R}~--6rQmTr*;Utsb`Kal1{9 z3v5x?J{Qk_E#kHYk87!nY`?%#g{@WGZj*E3J>;#41-5$F`c$;rqPJM=g)a6(^4MW762pA9k7piNB;|a=e0C6V z9fdtn`}~}!r4xxmf1QaTt_!i=)=k*D#_g8R%SQ3oyp6;mXE$NpiD3*qh#_ZBVu*_i z>qRVxrMDU6`2kj;ysg9M<)gNl+d*Qevv>UbI9MNIy>0cd^^MzYaxJjih4mvAgqf`M zkK1i>&cGIhZ9v>^xn)nu+_Kcn+)~Nkco+lwev!aZg>8n|2M5M&4PNII^4ccNRu9{t zxZNhl!3GN(B5WwJAeP>S2^&rfwTuupl31@;dK*OyIY$e-gIEwtZ)1p!h-@q|^fyk} zcw#{;z1>6}f5AHaTv$$@@!Si^8_AvfXl< z5 z!<-n;_RN-HUhX3n*a~5r6Sv#s{t!2pSYWG%ZC>1NljA1ZqOjc`x7%bx&Ig1&NDMjW z6ANPL?Lbz$=lwktKi>jd6t)d*>^3 zJgyXReJf_3^V7D;HrjrT+S76SA-PWM6VDI}Z1u2>ZD+U1ae*xg+p4(Taw}ftd_#Y$ zRa}wG_A;>`mfkAWvisz@M$T8_=UZTl z!d6nu=c{pBBROYai^BF=+-{Q%{jK%PJr2JvtnLNLIo-Hs!q#~?Eeqz_zOT@)6UM%e z5N6*$=EnuLC~WKFcAGr*z!rsVL)>nY4fFDbu#LomFq5@Sal7R?Hw)WBEC@4M+Zwmq z4H%-Lt5^--6!`R;u z_O7t^hy}6qw%0pe+fwxRe%yXb-VT8+3fmI${2nAE7TX=SAClv+-S-d+Z1u41joWQ< zTwsgBwyv|?CL8MfAb!3DwkT}d$m>JI?TgzQJm-hR0$V+7_t&;t9=AW9xqT$;V_}~N z`&8HgVnHmu9VCWr#Exe$`wl^W>@fQdK;Hse6t>UecFX7Nkg(5%eIe|yurGyuCG3c> zqr#368y8D&$Ax`O4Eydk!oC&u9Wm5$f*AVyUf2)Bf>?U{kr;CRBfhX5<`gzeep$+_OG!2gr$E$Uq&rqcmMA%OE3E^*#cV>wiIHhGpn#{p1EVs z?wR$2oH>XEwnEt6C}FqBbs{b&iG$@L7FepVZ6rB^goM}Jo7fM$zYDya?)Y+hW?4|n zg~Bc(7FclDeoEZ(S}rDWIA&c!47FS;>@s3OqK5EwE15Ioyqv_r@(9aI4E=p6w!;-9 z4*gvzET6Fa!U_nxO4!xHt`T;vuAbIA_SxL_<1hw2OtdwWqq1Sbt#ygbfrnNZ4Ru zLxc?#HcZ%XVIzc%B!7q(m29$|ZheIRU~XYO(SL(kmf{C;5{3Hw;sC&E4zc0kxcVV?;*B0Z z6GNRp5JR0m3j0ae&%%BY_N!;^d$dC5jB%}NMcc40Y$C~gk9{J zyU$hMcLy5LZap4Z;cwyO9`bd0&jL2#HIJ##dBWF=00eD=w^r zu#&=V7FJ4FX<=oAl@(S_Sb1R;gjEz)NmylJRfJU)R!!I~#PFO`omg75pVbgnGqRLu zUAsfhTSF`6EUzW<)D~7pSY2WDgw+?;z%%!GrlDsxytHVHjfe%2m9Jz!WW0WIb2cV% z$k{|#Q(?`BVY@UZmKN37LRd>-t%S7})<#%cVeN#q7uG>oM`4|a1<{;!CYBai7hzqA zA?Jj4mLOw)vEPvMHZRWIm%4f8#&!41jqBl=8`o1)VM1HjuDA zn#FGUeD))8X|{UU`U@K%Y@o0~!UhW)B5bIzVZw$B8{wI|zl`+E<|@r@!ZwOnkdRpH zlgjo3-wvZm9CF?vY>cq6!o~?3FKmLaiNYocyHnU?VN--n6*i3+w$XHAY0>tZA#A2+ z?*4L@XO<2*?-q8CuzQ8g5;j}deZ(*?bBLuy^DdqCKO!sdJCZik0Fv+^;% zhl!;{<69u?5n?zdAMR`k_~U&ci9>&jggq*3F)@s;+*j3VTV|%fenE7DVGaDqi=zD&k%vmKNn)OAL>Pwmobmz8zjCamcxjSP=F1v3T6B zCvnKRLD(C@HWEWEn~0@Fb#4~6Mc7th+k|ZwmL_b6us4Ze+rC8%^YRKg&LeIoiG%GT zmKM$D+rr)v_O7t^guO3px3E3LFuuK>wU5STcOQ6W%hK#7Z2LS*yag?tVjp^DixIb< zSP;c+5}#rIh{Pf1$HG1l_NlM~!VVI{7&eprEG-(}XCw}GNZ9AZu>HOuhMb3meJSiK zVMl}=6?Tjm=4@71n;?EnK2GA&qWSz<*f+wy751I56T-e1_Jgn=h5bYfb2eAZ+0P^{ zEt<1mg#Ahk{hcI+oWBYCUDzMOP6_){*k8i_7Is?L8DVFIofGyCG0gP?wQXbZ$HTu~ zocp)af1X*EG`q3Z$8LDeEW%QRWfhi9SaxAKgyj^LOV|a%atpgq*hRuF7IulSONCuV z49@|V6T^1TBP_45D}-GsET6Fa!U_nxO4!w&xsQiyJhMjt=H*&p*9j{q?0R9T!U_qy zff%-XVPO-*HTN4y9Ok8nu%g0>3A;&HabYEh1-22x_HiA%<;Sykd5mr@LHDR|1t1hgDu$sg$pS6f#K5rFPTUZ@o zb%oU<7TB&Hw)(_SX9Hmkg*6h^SXdKbO@%cR)?8Q%&)nmDOV4bsFfXlywI+srq75cCTl) z*WoxZFZvF^6k8ZIn0*&&Di((KvqYW|Bc9e(pImKJ|KC$j%Z{0a!-D%#s zbH!LDgzLJGOZyFD=^q)ZeSp`;JkiHMk!OsT=YBCR`!4cSD==ygh(6hOlVe==-L|Pn z65c;3`iy()&KK*B_11kzth>@%_hB)X)#19fjo5bZzh`{`$&(h1^$}u0WDAKQZV@re z<)g%qbFr``!X6`r{+1F$e}@ud$#{;jUonPdBo4Nm*w{p_SnP3Ofv_itp}!}EtssV4 zo+5^vD~Scs7@j7EoX-$L&Q-!z6GMMj)whn5w*%}NVQWOrXN8R^d47Kw!*e7KV}G6) z;$9$zF}z3&aW4@=T-U^lp5%E!+{+{m_6jlNd{x+B(cfz#?xU9H&nNn;kmY z*z3a930p6012JsdO3lu%6ZVD|=l-p<(KE|}xJ|-xiqBhYCUJo!4%-%Dn6s_Kf<#pK z-O--8W7zJQg10dz%>M{n@x@duvXd5A3JPh zyYmhGeMI_$eJtz~VV@Gid>$Z%{tgQJOxPjM+|R3g?wQR5=H&}v_mcTUorgu-m&6eF zm9QhkP-n$XR(0}Y47QuRXAX8$^ow}d4d??z8Cg`upfo}L=1E0zkkR^i&}mraTvod!iv_iEXn%>a{fxB?i6QPBF~t2N>|bISZwqp4M}O72hePnjfd2j?IbpHB$#L$v!z{v5g#B<~<~g&9 zxNO35UzD8Fe#5-rf9<5$!l=Rg^EKJMK5TWY%g)Kb{BtQeqCBim|6EE=(I?KOB%bcQ z+S&gJLOpgq#NXa9JO5$EYP0|2k!l>ZTw)yde?3yMVt9Xn7)Nel7ZMAiycZE06WPVY zu-{%H>{4Nu3A^fluiJ`yiiJ`w#Vu&k53~@IQLtJ5Eh`Ui(5n)Az6%%%o zu;Ri>5W_ydPkcvCNfCE5G0aOTV(70lF~pS-R#sR!VdaHY5LQuGC1I6?RS{N|80Ksg zc`boCt488rw+O2)tcI|f!fFx2m;l*G4$6(SXW}Gg+=d{q+@gyRd%3`U@KnSxU4$?GAMg3e7zaHc;dlENqCd zp`N*qonfB2wJr5t3%GH^NgU>PglBHtNYAW4#ElX*ni%Hr4r0hTM%Y+l$T?2fcwrNS zO(cdoClNy}cY5Z|=VZ^^`JCceVsqLGVVmk%yY$70+i9NJLex23*bHGah22FAW4~M2 zJ;YGwy~L1nmS^sK&i2e5`+c4z{v=DxSS&WjGdFH7G1M~8vpP{eySv{rTZXs?ggrmn3orcp}&`e; z_75@i_b)N@_a8CD#csF3B<~Z5%OWg=7;6qon`iDZEW2k`8EVNv4E^QwEb%r_ z3l3W@&ukfTULY*DunUD<b6zeikFdOvr9^eR zIj;=OeLZ)D$div4#*|-J0b-b^t2}dSx!N!8$EOTE8>|OS5#Or&)moOO`cg!jG;I& z)LBAUNnz!^_cy!ayV;9NtZD^?t(0fB40V&D^Rq@P?tLm9MzG}j5Ar@G0*s2r5*lP%@Ni0ZYjK%KBVn1Yz z-F`#PS|kp3E3v?W!&aLZ=DH3sGDvwjd)Ke7~C7&Y|QFf{j=*+ArJB&@M# zZf#9Gvl38SQ({3>V>4mRi6LhTVJ(TFzvUgRW4^6hiMXufy~3!awTR0_;!w*9QA-<= z6aBRn){YqVgZ9EY5JP_*g>@3vSy&fgT|INpOWfv}ZB5kKO;~qfJ%sfnhWU&WLoL0C zA+9$u#Pt!@ml$&1F07w0_Wo+@pZ!T3>Kq_!ATi|JLcW6$wo!by;UJO|wG1YPdD)sc zA)WkqhYb;NLxl|^hFXRT8{wII{2b|7xLf~!zaB*lb2ge7YPo|L;>Hleyo@D=oa2Zg z=WOwvf#bb6_wShro`pmI|JWxIL!H&UYvC3Ln?&L;hC78#7BlMhkFEJ5(J5(WY z=U_pC_xQ5Pvsl>B-)hg? z&*H2R_N=hyggr0p1z|4=dr8>K!d{6i#U>ZE08_YiU zlUdI7q!#Rx8-%?<4A)3E5(}(^ux;{eRQlq??PkwxVUUnmY%@8gpua674rACVY@4v{ z!qPnJ6V+&UJ3O;xsPj!>Z+YhC-07L+#Q1gzdz%<)sUcqfzC+^B-@C++^F3mSd!JZf z0}0zqvTYH!o5aEP5DO9+W3j!&5cdHw#O))7@qI`PIrkGo&SQz1lAjL|_YsMMeM}5F z&yl|sVV{sV^L!u?`vX+`-T|ez9ojZ?}#Dp zgs|_4A?FXokn=~+-1+>;Gn+xg{VePkVyNX;VJC^9zu$-<=kLP)5O#_f`YTI5hl&3F zByq66g#9h-G%?h2h8X%gE9@LGS#4v_D!tx5cLfDnW zQ0IN*ydP@GN8(`liD6y}2y0KCixGE~h`X8?YPp6Oa$ZXeW8Wm+w{RVaLw^N@T`w$^ z80st}>;}&oM*Fng752=Qq0SpUt7{mwBA(e|#1$1*OxR7tFuu;@`5}mE+3CH<*3lhD>Vs*3G z56Q2EF}~Nx=P*%AMUoR%i5U8ullX%_xj)#r#70RrpH~H-xhMM9&vR^ z9Qvy#tiG@Y!Wt4oosEd0md3)G5X1P2k$n(3o02$KGhxkzwGh@)SSw+ziD7(gh+zzE ziD9lclk3%pYe(W>?S*v^){z)$>ExOFcU))BVqrrqU4(TNb{jFwXE$P~bD=m7)}6$m z&K|;g5<@M$$#VcKPU6sCFJZljA!i?A$k~?|;{FrQ#kZ3<nq;~o8yIv-2gH_cz^(b92t{%*{Dd z*j>Wz7B*khd5?%YmLv1=trj`&6**@Kn=S0V$WpANsNop$Q;y{C$eZKkvDLBe$Z%bE zA8~C;WUPpmQ{G-7sPER{bA11gzX@PF}x}4En?_zr?6eZ z-WK)_G1U1kG1T&&u=j=S7Pf~NYS~K++vqRbe9rc(llVRT<^vJ8j~K@Hp|Jg)B~Bwn zeSYMbEk@3diD7)75JTLj!VU;KDC{$0K{SR##8AuU!oCo8SlE}sz7lqX80PaRF^u6D zF~l7whPbbVeM1a6zZLc!G4yvr*!RTH-w(o8k@p&-zaL2)>?dJA3;RXbufk3e!}xv^ z_B%0*;SbN+MfEGj zxGcg_gk>d$F=X=$%dLv+o)xo&QOn_(EkkF zONFWbiFFx?Lw}bOLw|XM%7ceQ8k>yT?a zvwaZ#T}v#m6~cC%uq(^i8vOBBP{ds?Y-iT3A;&Habnmf@`?9Mlpt}Kvy#GY7FJ4FX<=oEVGL!7q0Vx`$_uL?tfFVd(s#ea zZ6(i8qhV#wY#HjTBCM*gYQk<2R-G7XtwAiXV#8L`GsKwH^30ZDe76d#Ev$~Py2MaV zJz}V(J~6~K@T`#yB5Dmiv&G2SNLXWGO+0fyC(_h2x4&k>nhR?otR=C)It^PZVi4PGt%z$Uti3S*`&6uVjG+UGgLM?vNmyrLU4(TNcAKzn#4y+0iD7&_ zh#{^gF~r4%^%B-wSRY}1i3PS1!}eI>LPzqx6eJ`TyPd>geEo=FeEo?bZU8a#H;@=| z4iYw4*breug$*NyI)@WOog;*e6gEoOXkwVpJA{oPhB1sKhB1s2Hl7%2nLrFVClW*4 zBw~oWQ`lr-Q-nM!;>Tqb*>Qhl(3b;o+gH|KSK<)tP-|b z*cxHa3VY5o)M$^j=RLD!@E3?-3@;Kx+)Km|_cAfWy+RCeuL^rj*jiz)6GJWQh@qDC z!ZrwdgBbeTNDTdL@(eMytu}jR%MiDP7-F}2hUFH&%`^ACW7~zL3ELs8-KExR#_LwM zzc)qPTf%|J5+3434IZee?bEfsUUSHyiFY@e`AqQ4JC+VLy6iTN>MNUUZC1v4v5C{S;a% z7KZmfi9EZ!W37cD&xl<1TgI{7wUHujil}X|SKH5`wl$)UX`+t_ULU`RKJ5LrHbyp& z1HC-IMtRt}W4v`wiggEh>;5L{WB+>`+oz|fJ?{1KyQpWZx9%UJo|U4{Azq$SB9Fb_ zH?w_fwdkW#&g6ajPtr$Vsl)b{u)l?!_N;7ry2R}n&un3kkXY<2F}$8TC+r_#{|ft0 zSZsh5mHgO6Em??R3~!O`g18hC2g@of8!=ob$W9D7a|p{RESInggyj}?p|FdHVSE=8 z!}u;ChWFoIN(^zA5ySSoTv#4r$eCBz6~eA0hOy`K496|oGxK|9%TP-JVOI&eTG%zj z0_!wv*Ahc5*9j{q?0RCDmsDYehy{r%W3d~Eq0Yj>ZuD$ml+W&pcxKCxv#79Q!fqml zu@@Isf*5KkDePuprG%9xhC2UjWi!E#pJhlK>MSd)oUrnqxsR0!p1IqoqOeNBDig!l zvnKwLNb=ZWRYcr;@|ju8OI0t<9bYxi-0|H)47F4jR)ZL7sYwhugShq1$6kxXA?K~c zkh3;1#MKd2S6Dq_=!kP+eCak%z7Q_PE++k}; z40X05hVilgg@(A+Bo5Yw7;?5HhPZaZ+6(JI4E=Q^hWnyB`XYS*@t7kU^VVgFE-;_jqOtk@H?*vpjQ;A+tSmbKWOxj>SBna=H_4+hI-VwR%x8(EwFuQ({Y8O&;`LXL4skU;|UJ$kV*C<{T>-yIyUJ~n0&y`%m%VJ%-W|4|w!~0jnx^`V6@zffv z|EgHmu1mOiUyIgdb8XitGTZicUBbxFF)_J**H z!Zs1Z>xs?6wh$YcK7`m-VQt0d#l%p5BK}wA*IrKBjWC99gndg4{e4FaIZqJ7KEIB9h8A(( zlQ`H9#E^4BYa0<|Fcaj+5eiQaPG35M1 z*ePP@?@waL`4=(7{Y?yUr->o%jIgu9&I$WR*uTV3r}~`UeX!_q3Jy~G1s|8+{@y7mMLx1YCB-exiMh zg2WJaJu$?k5<^@eVK)#%&cdE0{=A8{(T$$jV)R#pSP;c6sckGk1JtJ#)ubPFQ(i6@*nJhVjj9W4$MD zzrYrSt&$h#{vBIc*tW|pIwPkWw_Df;!m5b=stT(n>=t3wg`E<0)(~+uh1C*vtFYR_ z>IkbVte&v?!WsxGnHut0>NIHxqHqg|+a^W&)2tydOXDvfdzU zuo2{ZBg{X)-BR@7zgIss(M$SGD-nnHIce0!&Pl`Axhj}{PP(8%PXe7$j`4upz|I-%w%0h@qC@!bS)i zDQuLm(Vn@l)9>)i_94_ch8V`)O`Mw@E8@lxLx1CiO(2F^CK5x=Ny6?FHd)vdVN-=o z6E&|(>%4g)?tWnp2zyZ2d|?lHX8RBxNA?=fe~kfSuiaqw-h<5I@VYL=R*4$o>@`;^ z7KZl^i`wix1DW-K&mX$4E71pg-2}t+cNWJ!7l=4}eUn)&3&!NSm>Fp zk7L*(VUG%1Ol+7X4%-r9=;+=z??qwtz5g%lj^QN{_cAez z;T2)83VTi1T4Aqy=8j>VXYO`bOJ0j(4C_f8Yy&aO^|@BoL&mu55-oZE>ZXBsiY?I4D@BjTF)n<7rVKjI6Vebfg zSJ->RFfZ>D!x(lGL)@@VwiobyP<;>b9+DIN?G^TcuzkWlB!)4Z?0fzgP|JQ22m46a z$HG1l_NlM~!VU`iOxPh|p9}j!*kNH`3j0df5n|YW-;3|OJxb!R{f-es&f~-o_cgJ= zcA>DnQ`v6$IP{l|7;~djwgyj`>g|I7yNg;n#+9s4bw*<4|K)gw!>g;9gm2rcmz z`)GMhk*AhtZILXzztywSj@1@cM_659^@P<|cn}vSW5@fGdlAEP zRBvH@h>fuc2wPubIF`&wJkewvLv8)UpJw6x?INxpG1S>#*Z^Szg$)um*t1?Xh^P(m z%oby=hY~~VFk(TJbGWd5iFK0a6LBL*9Oi5!G1NKAv&5;ZsFu;5*a%}$k~s7?nHX|TA%^kwD4qHErjj`1oJI__ zOeco88N?7blNjRCT4Y}5T_g@U?-q8CuzQK2mRZ8s_t&A8*(45jpJ$2F8#d&y&GF2Z zA#Sd)dBiY=Gl?%c<@@#hBo6&O5Lt>95jEs|Fto&BIa)qnU>JrN?}h6 zdxjWlSw#%BtQNM07;-)<>^Wi26GJU82zyc3OTu0j_KL7qg}o+hEiug5>%=hjb;8yQ z+aT->VH<^Q61LehyiT=;#1_wN8OFC&*fwIAv+cx?GfmhI&+1t*QG3%fTa27#2!rl}1zOdcG_7DqfkYO9$&Tf;RN3o6ek~qxg2g3FdL(UI{?I#vk zd13pA7;=7040V1Y>{DR}gdG(2nXp5|Fow^GVGLh*=I*?Ukq5<`DqdFHwceBWPt@`?$qD;L*tf(`OEFRla(+kRU?+$n=l8;X zAcmYj5<||Pg#AnmIe!uMtFV*8eiQaPvA{NI*lw(4x5?W!uti}jNepB7gXDyrB8EBp z(=+#R^p~)|J+u0d^E5Hkaz@x$VdsSXLkzY2ODwQr!#20J-6qd9>imzyA?GCWcNQ!* zIP;uYh#@Y87;4EXEE_T8%uWnBb9m-%_ne+3-Xdqi3R^DEY#DN1AS}1A3x!=o4C7

C%u%|m0p}{1+TD`GWL1b z@`<>@UYvU_Bfl4CJ)q9tN@m_)0WZ$2xHEXDuwQo)jg|^&cgc|p1J+i z6jn>vt-@*xt0SzgXEp>JliBww!0sg1SYiJ6D%6YmNO9%|>)Q7!q&my%i*@aLgWYY~ zK&)&3Bhk&

PL-(OzTu*!?7}t>>3q}l|2uAFm3Wn}8!C3P?FIe8) zLd@-pxb5xFz&b?v6;2Az@yL*C@Yq$ajnK zOND&*DBmRHdqnx)LwS~`FD@B)?}+~vxaHXsH4Z$UGk*=-@)HV8qrmg+Pm%v6aLZ3B z@(lyeFHeg6&w*Q>c5M*&kb>FUPk~#WakNC>`IfCX=f}V;PunjZc)U*h5V++Tzl#M< ztY+f-z%9?Q`hgRx@$Uk+Jjdz+~4S>!qPcgPc~@h^%z$NmcWy5^`Z$CnwOo7vG7LtL?rajzJ4rC7d& zFc4B|GrvVpi>g@f;<(b;G3Jb>iJx%{D!7-gsEui6Lci3@M9aNLd_1%HkMOW--jRA0e%g_!+munhA^BA!Tt4DT`xBSsX*k z;uumE`?;3mXUbwfQx^M~ve?g*#eTx#y*K-ZwR1ni=gqZMUujnO-MGlx8|T&pX6qWM zRmcw~6W%Rv>lUeH$j3ge8}h6PcFsB>Pr2lo^DW;Z@VI@OhkQJr$MxMhVX-b{aeb#O zuJ4q^v7{`nZ&+O4?8(;4{amyjGUvqg&AID`bL0AM7xHm^TRyJuwjm$;*e2xT`raty z>KvX@7TA_As_o@`PjEkAs_qJG32pjrEDDXv2T`-ed`eN zv2X1|KJ|@$+9j~qx6ML6_RaFKZ<~gE?Asv2Q&?KK89g$j82I74orfmXCegGUQ|5wg~yuH`b1UfyKTJ2>IAI%g4U; z5Bb=)ejy+G);Hv1-z*>d)+gj+-+G69>KkkGHi5;y4Ga0$H_OMq4GsC&w;>@P`!+b_ zW8W+v`!*=#W8bz8`P4V;4G%2#ZQGEKeY1S*+nA7#eH$I}*s@Yag?#Lr%ZxV{f0<+fL!!*f-0^zU>(Dv2QzseC*ryAs_o@`PjGZLO%9wY{;j+vCgOCboX#> z?3?9d-*yZ6*tc;ZAN#gz$j81}KK5;wkdJ-aIpkB{*pp8VEN;_DAs_o@`PjETLq7Ix zkC2ajn;7!3ZmYOVlsty;UFbz@dnGa1Va}Z@7g15(T=IvCb7`l;igU>yF3zP5k0{P1e`Ikk{lV&!xQ_~Kga6e-1CZasiSj?b7Oy)^UfEH`TYXHsIN>gY+hKfC9RS|vi`MC{#&IoIoylt z5Pm0)Ng;3j4edoiYyFMGigOtQ z@z}pQuz2iW8}jkkzoyVqPCWLn56tSFI=HSlmp&Jd{TmWH*1p7fQ(*DfzcD$tLpYay zP1{npHy7tpW;peK?R<-1__cu7jJ|WbU<<9#knRv{*IHD$Q?RUHcL|2gyB&6q zVAy;(uaNRSiR+nrg^u#>6KqQ1=lz1Ad%$52I&7iC9&*^j4tvC5k2>rzhdnMB{pAV4 zXon{Sqt2dk*wYSs#$nGo>^Z?0gU<`bSb0G(Vt7$7VqZ@5L0CJRbdr5MzPu!CQr^oB zd&Oa|3dTPFHNhzFb-@@9ZwSUb@oWAeQF;8q-V{3e(v18tLt(dd+u0@rFqTMO)J5e6&d%>6o ze-I4akAgA3{3KXb_<4V8`=wlGu%CsF@_rEv-LHbNU;E8rzYE6m*?%}p^?CgKQ|R#X zL}xwvOX%?BZ^5v6cHc@!x;{IJ{=%~ZGlY&k?LSU={|ZLz{|Ux5Q|07pKVkI*!)ATK zh<$GTYMZdtMP2V^r-$@KzMz%sjPe!}<-rygY+S9$$`XQQ1$$o9HSDzkc9Qm`fn&3w zV7nIO{di2>HX8{Y@ilhXk_l^Dlx1&C5@yG+xs=0}7K|7sH83lS>Ls5R?NCU|2px0J z2=TlfY+0eh&*dDpyu(%yjQ!Gzf>B;m!B_*---e)DN$6MuRu&B1DuQKsZHDJ71Vgu) zV2q>H1w*%nV6?-7;yYZ?ttoW)xt3t)nmKH3!LV5)812`6m{riCK4>{&Z!UEB(n2uy zVCwHC(XAtNuyq~QQn1XN4XKr2*laBreel-NrY^@0Q-8++TTj@ey!9QnfneBF|MMQU zp`+VKFnnp_u(pCxp8AeVY_@ZB?FB>E!C@OatfOG~ss3*U<#lp&|Kyth-R_+oT^GS9 zudBnlIjp;2#HYq4{nA#EWTx-A7GhOGq4a*N?PTEVc{(_y_F z)>|-q=_A;T!k4~+;b%X=@Uy>Q*c>3(^rAdni z?jRWLx1+;$5)5B<77Uxa2!?J~!O%6yZ*Y|R_3TP*8t3SC6O8h97mRy{y~Vu?*my@b zK`_djC>YnUnfRUw+Gr0)r`9n(E8J7)@MV(2COd2|hfNWTz2X#cP09V;sX~XJ(;T+9 zVC+Bl5clL@(}j-mW;kr7U|f%R*CewNUE8AH+FQS^?WgY7vmM=>gqam$nCq~4g7G=< ze8Faz;gI$bjPmvsEX!4u%6<;JtiGL8j-7j>`wJcH0KxF(K!+XVu!99-e)&}F`6%xY zp@SXju)`d7xWkTc*pUu9N-*ZuqaC(aeiN$PM$Ci92pzE>D;VRwv*`2ajuSfA@eW(y zuoE11V#4CFd6L6UcGxKnJ2hc80*LQ4hn?=QGaPoN!_E?n`D}Ob`51nlEp)JR9Coh5 z&J&FBe!jyl5DZ^3hh6BfiyU^bU|b_yA{f40Dj2%U1k0>rgmk%J*u27FR|+ z!Lou~Eg07p>N7XykhjF=E3j*XP1v=95#MzVyWU|pIP6A;-Q=*F9d?VuZWWAui25vx z*l!a$*zJOiE86eGykg7qHM%;Z>8 zm@vj(UT(tT@*YZ9?B~M{d&FUnI_$B8n40eHKhz|6%i4VH5V8VEGUUe)bq+-j-u$Z`0Sg zKk>Y#+9(qdwmk44WSa#v1UU!v;Hj@FSt4ypIK=yfJy_ zUsNygT=$94Ve?ameI^*6=Xf!EE_B?F{32nopI;`-S{gsU5{&WwwP4J3-u?1#gl@hS z8dA-+_H$8vS`J^n6*|iMPB3)e3&uR?T|50CblCh+FxJ4I1Vi_;!+vqtuMYdoVZRGT z-~B_daaQ3WsryUB@Tbth{t_(9EmkUj3&!=)KM9Nb=f4TdKPD|~{+BR2j4zc_EVrmH z#Ww3XtbW2`n~OPY@r1?YEs-#@$$iBJf)PVQ!I%pg35Kq*!#*emZM%)61vPz z2x(=((5>RIRRv>BRdXwBHKC)t)g88mV9YOTI&3Y$D6g4dw8Pp7i^o`v!X`Cc zIBcDS<(K6q4r$$l*)eRkbXY6F=z}fA^F+*%t%Z*Eo72)tFZVNS9=o8L&9p-yttV{G zs700a1tY!<9JZlgSz&V{!LZrJVQn4O&SC8ZW8Uo`7`|-mu#OJv*BDk4(sNy z?t;+{n+V2p;@^tT23d196*~4{n+ZmJZY~&muq_0m?|wDZI!?J?V{=QPBfhO1)>>Sr z=C|cTNj-!P*3)6?|BRuVm~X;$484R+=Fi@O5qqD6&9rJNQs0ExVRZcx)(cU1@1L+( zH^5;79k#V#%vpm3!_UEj(FeU}HHRd+wo!aT6Bfre%wgL&Y`DWl2u8o|cd(US?h~0E z3TdR!Q7@we<8!rn&K_m7(9NtZt};e2%G*{jbYlfWx1C_<&J_Qv5Z(4d2iw76I|{~q z_W9dbpyj%z9d;5r%G=ptyEtrDhm8{~zfBiBoRj~tYuQhH*-hveM<3*ylCloA__*rW zvAMgj$=owuFmw|fHql{wIBZXcO>)>|!Sawp*fWOL9s6EF$3E+bymKt7&v*`*B6QTt z)P%*lX$dpW(CsZ4zDySk-3-BIn&FUUI&7B1nu|Ka=4_$EmpOv5k5v5=-CUuA&2!j% zhwbCAeI2%+!}fRB0fMo|RL>z1-+@90JIG-N3&wo*_RiM&i~6qR=t~Fnu4deK`{wv+ zhA)SR@@CrUAzeApewNEacc{?u%;jN%p*vhK;yXexbVmxt_Xiv$7`me!c8tT06^!}Q z`@G{gp~IKs1si8o9MS^8(4FA06B8ECOD84F%E9Kzf-$B~5sZFzs$g7Ocx(S@LdP0@ zx?t>k+U7R|bss!K=;%x92p#ts);rm1Lfbr3*krvu%VB2=MqQsH81<>HEzq5t=;F3L zFJW=po}aK-cR|8Tm)VbyGQsfWLWf=Cu!|jbiNh{+*kyucc?97Z5W%<>y~1Hv3P$W# z35LzR#QqcA)j|im#$neAMts)^hRy2*n_1M$I$|z|-5_+>yiqV>zezB3Hw(u6a*JT- zZWWBrxo#5--R%y$Lol8tR_({N{hdMwyGt;9x!Ym)IP6}*X!rXBqrCea_JCm6d{8jf zxP^jI-a`(1IAQVJ^GL$1Yhv?Jhdm}3ed*ZNW<}5Cj|&}r=?TH`&DS0FhG3NUro-NH*xQ2fJp}IvMtSc#>^;F~ z+xG=y-}8ZBl=q=v%po5MhVEmBeUh+v4*4`;!;8LYZ=WU1j%D^Eq|XJzmoFUlrNh2* z*w=#b9GSXSqo3U!3cP3*mieVxQ+?B}|I zp=&7^@wE~xGs7XZ77X2bg0bKA+F^a6!{!DK+fXp}ZW{^4ekl|8tYAIzNufMn<4YT1 zb6jn;R62>zr(ta!T|2?D!e)EH@TG&pHcnXFKRYJOio#~6gvGkf2^&%PXm4E-X2+)5 zkC0{^WKAkAQ;!aKMb}& z%f7 a_!sWuRkoYr)VB5{&W&J8Xzx*c|GxVS-WKHVzvu80C#{*hs-LtHzK<3C3J7 z+F@fHHq_Z)Z|msBI&3?^7}MK3>}+v=mpa?Q(e3E4odm5QZQ_e z6O3nVcM}ZV?hYI8unB_kyv9U_&3FEeVGp4r_B|aoNiePhCJRRFdpT^1VE8gsFlK~c z_>N9kzBCs$k4ac;^H{-HSAG}YP00LmoY3Lt@q#h8E)cAy@bd)0xR-vSVEN-DL3@(J zP8JNCrwGP%)v1CJ`)Pu4@8Wd9sFyPYLwBZN=*|*sR#9hX3q}m*2*!QWa|J_po?x?z z^3G2f{u{m^Va#*JvV_qF#x6`4aT~j+VD&8%Qp_$1jOQ$jU7Xm7eY@0QmnCJe$1|H( zBxPKloZH7vD$f9TWbE$uP92WU>Q{v-|A^(4ApL&rE_MbNg ztzEA)E?9@?oLd6hILa@7(*OVeA0A&zSzS_}w-$aAll2?=&TR>^enMZlJz>@_>}BbW zg3&%PyE8BwSB$N@5<9VPcRTE!qzr2}+U~xjjC+%F`-F3e|Ng+@GN})Iw{b;%JP^)R z^~1R~uBe{}!?~)yIJZYQm-<>5&W-Dc`lHXr=ROq9wQ)uLJ)CIks$a^DV}B%Q`xv;iZs|=Z6;yTXB7`Kj41~#*F(@FFU$FkF&zd z_fr@HuLvEi-tpDC7DXJd3LSnn5<1wjg0VkXHO!5vpNy%Mg^t)?6TT2b`ai*~1j>6| z=wNRojIvDsX2NFK!6Kz>cR+YgS$~I(bXaeP8RBn--b(y5#vWY74)bDvJFywZp0YUh zl*O^9ERH>8aqKCJ+bHeVv7af6{Y+WxXUbwfQx^M~GHVC=*)~a^k7Ib}R6BFgxQ$~- zSsX*k;uumE$B?o(hLo9STzjQsI*#F8Cx(>8F{CVxA!Tt4DT`xBS?uR~j-M%u{Y+Wx zXUbwfQx^M~GOKImv-cA};}||j*enwkDP?gCDT`xBSsX*k;uun9Zs5;{iJ!)Z&l?XP zB|7YxK4o!yDU0Jv*+?^1q?E<+r7UiTj~zc#7WH zv7af6{Y+WxXUbwfKXv>}S?p)ZVn0(B`r!~PHqKka`r$7{FcVEmuvxDKhOKZQ-$ zUxHCjN(WOq>P+ciN=Lu2|7o4Z2Rl?S`t08#2G~D>(SH95MtT1UhRw=p)qc_r^#r4? z>kEc2iwQCEp1f#sA z1jEmz1)E)4WtC+Fqpp_~jPjNfjJjT4Fk)XpFl??U7-OZW!&VZEYqXUG!_QR&qa9Wi zjCNQ}Fm&qwCC@JEdUc^AzBLLK&(YM&nmJpvo~ZIUcP&u{{Xpqxmu5nT&9wz1-WtKE zr{;p;OAEp1cS<+As6*u^?YWMyNen7>>R$DcO^Wi?6*jql)>5#n(6thbd!6e4(@rmR za~jw$i>^s6m)VbyS__+5PAZjqR<&Qsx`hQ>&(ZZ;zHZ(6j&1{oZRoIOjjGGb3O_e; zbZs2g)?w`&HglP}eQ7Ur^xY1EF^<&Vg%MwOr@W1YO>`X{)+gUQm)jj(C!wQHbQX+p z+eI*Jb`=bp-2_9|-C>(JY*WFg>&*nCyv+q;F7UoLe+!`_zAXi#ysaG8!(lxg)=Mya z=`9$(^l?~U!LZrSVf`I8Krnn6C>Z5!?XW=(8!Q<6wY{DBbBNI4=TL_YbJ#Wx8}6_X zf>AFc1;fu#f}tBN7=Lekl~XULk+8N_6-C-JVRo4QGJSyM%I%lgp^zpCU6zyZy|XRs z7pJuUv;I4g8}37fR-6v5C<6^!vdO)zwO3r2a<1w%K(VKW7rUBq{= zv*+6)cUaf;EMb$q{%paBVUA$Lu$5EZTt_$0Ve5-MG}j389o;^H5yQTM5yP9JJlNtZ zSf!Wi6Px=9o3Q-_!{*t}ymWx0J5VsnJ4i6fJJ?}|IBcM}2Eyh?Yldp6qh1abHu3W? z!LWI_!~V=C*>XFeJHpW&DHy&S<*=g#qr77TqkkSN7`o#G!&Hm*Yp~L32f?@MI zhh6Wm8yt4p8kLaBa~8hb=;&^8*v*1r^A^E~?^cK1=CIoZqr5u=n{6S7bf;ihPAZkV z9Co*0ly{HA?iGyk?h}k>WbPM?@*Z&5gMv}sLWeb5%POJVC(u3Q=pJ_1BM$pp%x9GM zsH1z#VUIiPcW3Q*Lg=XLCk4aLryTaQ!=7>2mttQ>dCxk!=LDNosZGy2?1hB&DKz%> zV!~=7_EN$OWd(oPVXrvsRfoOiu-6^-hQr=;*jo;JTQK^|I}UqSFzW0*hrREx4;;3c z)8{{QbRP*ud>;$OKKT>Di0_e=taIr3<%Q|2L@=dkY`wt;xIo7jJFbUzA4yZ_{{pA#0(FTW&ggmr@={hF{?_glhD zk`=n&1;fuj9QJ3zh8Oni?XQH{F>L-_u=;keNU(nbi=PGiSJ=r4&3_K7oNg!Su~W}s zjn^&f+7@NmTYaJ1wcv{hmKAJqhutgo5UfE=#|vO9~ynG;!F^tJQ6DDM$Cm%603Oc68?~Q@3szN4KoQmUGz3%hp}q@{Vo=!I-~S z6fCo1LTV})x|IaWa#f}BjI*~|S?E}2RuK%Ft2%5ohpjFcF{~jNp9iif7`_b89WM7> zK06pU)GDFO(5>azY$h0ft}PfgPs<(F<<&U4=7LdP3&ALF9fv(GKKDYmuA^(|uzO}# zLMq3HpRF9-yG|c$Ep*Ju>j_4A>kGzxoDC8duK^n-%p%I{M@SnbEM9-xB+SmIytaa6 zc0x#-9cw?!b(YzoklG0yHaE-vIIgTC_Vz-Tns@i zgf4;+!-dXgO;4#+rfen z`w)i>6^!_XIcyujsI%dMVRMATMmp?*q&`Cv{crvpC3M6-S}=SW9PuC+?9`FFOexY-ho!vt0y3x2s_2#tDX>>bWay?k04w z-36n(Go8;I#yh$Rg7LY+M8PQUjO6}LJl>b+QVB`7?H-QJJsmblFzRKpVEDOL!eT$C zB+R^K?U*VUHm5mkZ--5H*bKo~!)FSHF9(UwTVN}SYbWY#mSc0aVCd!uMtO4`HqT*i zi!~gZ^Mwvy_7M!-z7E^ZVf#Dm&ZN&*LZZ9_9NmF}5yL@(Ve?>zHLMAR*D=PHYf(2l z#Ibp(!wz%U;ezqp@ezV$=3Yog3Wn||!LnRcsa&$6{ZgJ^#ue;nM|X^1w8OCuJI-Op z3zikWEO6KvE7?iq7|@;I=uULlNe(;NVIM7Bw=btSy1ksw|4(&vmy7?gP3)&Rx=G@B zEp(>~9c_DtVCc>i4Bc6Rp?gDoZydU_9o;z&JJ(_7IqZCgT_6~KW`YslmH8hM>;8G6 z&`qq=ri%o_<`{9ki7yum9qbapuz9Iq=q?is-Q^CuLNIJz>9DH=qr7v)S`52d=q6NZ z(=~#jJ2hF41{WH8yH@DN75vvu_KR-Yr^UUetk7L2Y{IU0*bNT5(P1|Umb(}HxmhsU z=oW|FDi|@`=CH5D{ug38PoSWx(6I~pjhAWB?3_rh2SbXo|tAtfT;`5xZ z1;dwb9Cnto=lj;tedn-e#orlljdwtA^IFd#-wT`6^$&s(-;WM^MC@(YPwe8%dpH4+S6W5LiZ>9DVz>;5K=?vj-&A?f*LDMxpH{)HB09WgBJ=ynj}miU%& zbju1xdCLig&E*ASAGv~H=vGYF*icKPrU|P7g!h#aW+?xdD!i{OSY{`Lw2ENFx2j-S zuBud46O6so>Vlzrx`i1iw;yZ`N4KWK_7U?0F|6h2-VoP~=$Z)~Y;D2DSy&;}2$tof zQfcn67J^}O9l@~qp%{PY)=hNr9;Ic%%tdUra#(A@cwTBf!LWIr^DNv2&fex1aUYX* zSYMO}+d!~=EXa^H6fDb0c>YB&o@Hz!7`g|<8jde*9bG$zZR>pY(%#W^aM+Zl7N{P7 z8#}taMf(w7M@RR_G1cXv>m+or&Vo_bT^#n9*k@t0tI)x^2}XI{9kyVp>iC#@t9JZyywi1l;dI&~&Jss9dFl_b~44Zuf<8Pz-3Wm+0&RpKl(d{C} zEjIf*x&eX_!$85Xd8hD`Yp$(@j(g361f#scf>GYU{6lj+-iHVsD{#+jf<2PKU8((u6vqyHN_ zBw?}NhbC;835s-B!bSqZdsPule8 zz~+a%y(}G**s;Dwe>*m@V||u>dwgK=xyL1TtY6YM7bML3C4CcL>xE;MtY6YMPe^oe z9498s`X%*tQo^iXQg0_G%=#tuc1pphH)~(&?bN`=hJ0Q1W_^!(J1wzeeVBSXBe3|~ z(~EP7zh^j?dOI^=)=#N7d_`ydlzKZW(Zz9`ov&4m)2k#>EAK z(T15|%=s4zhRurvV~x33Fg`E3L@?U(Qo-m?>VF-fgDqi|98%-(s5F!)6!3&}}Xl<8+{4)VccK4Se6MVwAVCV{^D*T!$%}j2RUlV|22xNqkp{ zcHpy=&v^khVJHq z#lF#BZwai687k7P!VYEJCKzShE*NFpAsGGsPQkFL{vQKO{XYiUSLt9%ho5%|Uto6& z#&}WX!BlxIin_i>*rcxOpJpeOuV3j;_X-{La-U$taKB*aRC&~yVvMr~5}R>felTJ2 zSXx-HxK9$_!#P_tULH#9Sj|vRj|fIgj|xUzJeIK7x72Sd51WrCy4aT|1f#qs1)~n1 z5{$ZkS}=6a2u8g;D;U>w&k4q!?0LZ$tE-7ViJvbB9d&lKn5P*7FA5zosPeF>{&xdo z?j>Op_Of7<*IdL$pS(=OfG@8Io3#6@f}vC8QJ;#{SZ53AHDR+RCzZt2nQyW}_omRn-V%)GOS`r<1LZkv*TQDcHr0%EeL$;fhVJE-)r@;T&kLK>^_rVi z>+o}%eDG^O-xl$)H)-U=cb15O^3K?>+Lz1@h4hXnkJ#T8jD6sB?M+>dA+tjvy(e`1 zjoo{qKB=>1MGSLmO;*+sj4?Jhxud#;F@yxuGM~`>nrMWuALsz?ZTHVCzZ;lA_m5EW8o*dreY3RSlIkb z*hKfa!(QFU43ygs-4{Z~7~5CW8M-fp4)&G9zINC*g3-^u6^!z}6Krm!Ha#Qe9&COu zbgab}ZEE^*opCSl648FN(GS8Ved!R#muCedh98AZ=9g`RFO)ZL%jy`2{cfl4{v^tq zUR#}&4sEM#68q0W$2I-)BEGD!`HRr8j_n}gqkr}k>k+!OyVmXJucADB`AsnDbC~dj zHSS`;sF&Y`&ACPF-Pf=76JOq&SU3Aa*hKfIU{lOgNWY18pwEvGEX!5lJKTg#bbp9> z0``+*^KW63^8OKQeBnzsQ66*ohN2xxT+CsM3x>@l1RH0rhqS!dw`Mu1R2m2!*A^E$@ii2>xrNQc*0W!9 z8#NNTtk5-f*ph;=XIMk@7wV;n&@qnQ-l*D7*hsN6N=coi8@2K zjL=bM%Q|d1!RYg9Pm9g%#hQao%@gRB7v;fLaM+3tTRLx$a$llenmW3b1mpU>qv$X6 zr6t6?1p8d*h;L<49`|1=!Y00Sacr(4Y_c{!EygDE^~r52A(iV4U1QO<#IUL;kGft> zFy_J41tY%qMC_EehS0&*6pTK2v9o_!OX#rK%wc0Uu7p&svx#;nq@zR(#J;p>Kju$$ zFM++rW{%BI#B=@BXH!ugevT0=%Oj{%)^=jA5sZ0oZ!s_Nx#5SRJYsJk;$wZEAg+(- zpGS&0nLgND_z7!~u)K>Dv9FV`p2cx{TQ^~Lj2I5+X1Q{op#6Rl@zE!mh_!4y!8ZQuInYrL$|)r(O)(Y4Bdu;q1#9>bZs2gRxtX+fZmmm%KeNwYbM4RZQD-R zWG+xPsp|<&I~?w?_M$xI$y5%t+Y*o1AIuy{;&bXX^cbxxSoJav7Y7!UZ; zMd)B%6Bhf~En!w3y6z6!Bw?}5O%rA|C)kgWMs~8F<-RoD4uv#Y%suppdqwQ@&%;kJ zb=l@rI~3Ap!WYU@_ic!MbD`t2kgY|Z$IlL;pK)z}fVhsMZGRVika_S(5j%6M+Rww( zSRn>=Uy(Xf_X~(`O%Wes`g*~p)HXv!)fw?M6xSqKq1!^V1MAVZ-R-1uyQAAu=$M!K zi}jCo-%9AH&woYx(YC*eHiA7O+7>^1i1P5Ir(m?h52CJN-gQ7PVUu=H^C0YI(dUVw zx3J0BRM)=vsqWj*Usi1&nzoL%?Z0X@!{(k&+x8KD5<_3XXor4+F;@Bu#-3=mq|V|w za)8iL-ax^avknsNP7GTM9c++b^y^cEP4<|lihj*?;~=4WsNyjL;GLSwaU> zb24kdHqQUX-By%GosAU?KX>YE(U$8PU+xm^PQ7d=Y%(7H6l09?wi9fexe(Gp9rK?D z*2{CXKeFQY?L~Ra*Kc*TlXMI_2pw~X`ka~1i*FUNk1u>tpRX{c|8~~fF3$R}qwq5; z%G*gW>ho;Tm(cAjboAX1MSSRX5jxl-;$8x5nAj7s@A0k!b`|AO*Xr7r{y9$Q=o81a zsf1MS6ZpBC(9N^cLsIvQSsQm3I>y6z!59zfxf$Y9pQAG#)H4gj-cF1oe3>A8$@1vJ zGxmZp9>zL-Y4XNqpj@ByLG_s_F}x`Jef zd#nuxUuo9{X{_YpQJuV(v7Nagt& zo0kZi)XTn!&3K=@U&5>e+WiZ$*Jqwkb1P#--Fu>5J{6yZ-MqQsE7&cE7jPbC1w`xCG8#{;?DDNa;6F=vQ`4ipALI*oVFy@Js#dv2O`${mz z%4o-zQ$=~iaGJx~3Y*mR=|V?*X9&hVOMTvt?o6ShUiK3GntIv9SvSuTHd#B)77Rbn z5sZHA{cVJL&T(?3HmT2n*!O%P+Ln4bSNKBg>ORr9qCD?evGasYV!ze#<$R%|5BBe7 z1()XnVz@x)SdY}_GUzg)gIy>X*Ldpl9Oi-@oqm0hu*qDfKJUkui-it$iNh{+*kyuo zP4b=d`O)P<$NI3FGuK@qboBXuMV%4fl|sjQdxTh57=P+BFvg1d%$xGA66LYSTthtn zGNDqNc6Va9TG(WcyhbqQ!D|I$Y>rQ2kN42mCAvO^kM?$b!t5AzHfIORmB$$CwE7#I z%uWyKfkyVTtYfUG`})LpgYbnp>$*IkvQ6r2gj3fy3Y(d^7t&1*o4b{rQ7#WZZx%Xy zxkWJg=cZz=!{)6*2fIzMapqn~w+qG`^7;{WMmYxN&uQX51LfTzY-VOEq&ppUmtfeu zTd;|>CM(ZuWxtgDM0bzSWp;W<_XP0={NnI}#d*Lkiq*A$0_yW7% zVGjt#9HRE@*nCjvU<(DKuGMFe)b&F`hs}o__K0A_pq_KY=A%Lfd(2^v3&wszU9VvC z388~MDHvn=DZ$V^Eg1Frj9}=V6^y>C)=lbD{kdvivsT z-CmX~KRmqK%aY|sgm-&cvi!*KZZAug9~Iv1Wy$iR!@IpKS$<4-x0fZ$ZyVn2Wy$hm z!@IpKS$@0lZZAugx3;7GEI%&f?PbaGHl`WZmbd*cdtl3t4|#i8vb>E+#--(LZDoD6 zyp46nyX9?6F|I6cpR3u+lI86>#9o#xZ}Tg2tmW-GjB7E=+gfKYOO~G&-tA?{@-`1K z7g>IK$lJ@3r>0y{K_0_dD|zkSF*faJKM{WZJYWaB~Z!b%h zw>>=jc+1=MA=ikOPseYK9SZ4F!Dz?N9QL`xz7UMM`cg3J`76P8u648WwP3WzH-e%2 zRxrx@PB8lD_kvL`KM1x-;pdNn5&KVq5yQ^~t6$i$H~Pyjfzem&_^*zg-vlGx-vy)3 z{~;LT=TC?I<*>g6V@&@e81uouf>HPXIjnM))y_c&)>VhhhjrB<5tgh*^r<`G>F4{i~4LSOK<3!gdFxG)36V@v*#%Gg+jS7tUbE$;IHq+S0 z2HpJdvXrvn;TU~_@6M!N;_{YG%Cmg}Ytb?ZvvQaZmKBUOZaKkNuhjQbGRBt|y7^XU zNGk|Nd@BmZ`rTA8%3Db=+Gu6LXxmi;qfe}wuq^k!_Q&N#UuqB?V-8y_=o&}I?c61k z{L)FjCdn_1@|Pw3*3kTdub%iAGw*vR*AP0!$(jX==P%e=fpss+wm0>if;9zimXy)H zIBsuiC(MpfMvY+XL7EFjTec7ko9hTh-&t2MY_=4v#+(hQRl!s}wGOOP?P>YXw4ORU z%n;?Tmy{pZVQSNK^X*4S>nA!p2K)C^%N=y!V)hnNo7%Vh=LSL7Hp#b3^6?yMh_=~K zlt~?{?-?Y<^m_))ChVw_EVpRRwp@+<2x+6Fytr=KB+SmIFSHem7}^PjuDxKy*Fi91 z*jO-Zc63-L!LZrcVO<2Hysm;_vzuV-C%X&A`mu>%?0YvA3}5&Tz(&?MMXD*@Hw(IW z4&5T;<2lsw@f>P-v}VHcaX;QXRP+=3@Fv0%U&-zr`M7NF5QD+YguGYbN*~LaO{@6!pIKg7*rN)Y%*kq9q5UQZ9c${K{N{bRZ7FY}(82adSUyw>Kle;n z+v2#rO-h&@n{PitI#k3?J6tT*Kl;mLVUzWCFTq&Dj}XseXSpZgp19D_U!D>3C%UOZ zN86q*`aE;%N1_k1R!tK&iG6Ru@MXHgW(Y=I&ve+VgvISXJ7HFp_%g>~a|Oefd4gec zzQguOSZi~pNc$$t4znNJPq3`e?JpR8;vumYrri${44VfCn~cp%#XcF`fkMaULh2cI z#=}8Ehs}cpW6l~W_I2nE5jy(rp@PwO4-*WVhYN-;JBjhfJ^P7b4WKVI7kh8YJ9e8& zNacBux%>#>3%ViVZvs|th>Ux2&nc3+fo#3z&1*2d8luyFt^4J%sdZDi8i5R$MK1q~E3?~bg z(fPfnVlZe8G_M2&lHTk+W`wqU9K0_m1|C@o2}5dnq_t7I@VnDGi+Wibo7a_`7h-@isb=8L|BFCQFV9Rq#vC(%aK=k>x*e7Qj|_NF%qMtL_0M!&vUuyN*k zNH2?PSlU*#ANBc}=-2emTSR%VTLq&}td-w5FZYQo_p(yCP3VaI_Jp-9H1>8!!t7XP zKSH`wuxYiZqW<2IIZJ)FdRFM}5;lqNGG~n4Ep*ttM=-|9y@H{;PcZId-Y*#a`Ux># zWAg!_gFPr1eQ=W4>(j5F5MzaPWudSMt1reh>)1m=M}L_t_AkWnu+VXBQ8Ts@Qn`P^ z?i2S$@#PUQ#;EH@M0wQNqk=J)KPDJ9=ZN~G5AG-WAoaP7*t=0LkBjoS2mj@W>UyEB zy=&hmgiXqOQZVMFrAAhlhwdq%W3Fq^uUg06X0Djeu=%vGNxQcc*Lc+D?7`LLQ7=Ca zshj=dw8JyP7yNuyFy`Io64t|tFVgb~v%}O`?G46qpC`V{TiKaqmQUf~=nKMTjh!CS zi-OU2UlNQMUM^VYS~I0*Azlg0Vj?bkezj-F+sl&Wdxdv8@mnF+vJJObxw?0Ncv%HNB#)svtPtY$c zZ}Ta0s^zU8+RKvVZA>yQEpO`!YmDV>kHtRA^0r53pKkfdA#X2BmbbN%_0sY-_u9*n zREI%i_ z+sl&W?V8wLmMlLnyxYr?<>!ZYds(vlKH=S7mMmY_Sbp`aYDT@iCKxfjE*SOshG6uK zHy!qtVAS#3g5k?Mf}wj?uuUxFklquFHhy0)`o{-?)#Mf{l@A3Y_KyUkk9{oI|Hgx0 z^d);1JZ)RX*{4EBe4h!1pPvgx-~7U1UkZlJuLL8$uO0S{VEFQ_V6^*pf>AHuCoG<; zen?n6FVs(1JO}ZtaD$>B*jr8U`;S4_ILiMV@=GT9rIUP3l3y6*e@gtYI;Ng~NmxAZ z{F<o=e+$Ms^N(Q6Nh=A4FaHY0TJWD>^jj4lYlLEq^~%{6 zNqH|7QJh}AELZ5_`o+)X z1MB*Kw6j8Cr(twHzWx8G-?VR0_KHcpn0v6Mg4GoMujH_m1;ZxKQ!|If?Y@f8QQoS8 z(eA4WMtQ3{Yz@I^+cgEl=2{MG=CHLLRwEeWtGQtK(n7FJiZ)tDFzRz%!I<}33da1_ zN-)NG>w?8&h4R)5%-Wi?esQi^i#H4NtF6VPEkfSbVv^-;EhbsMhD><3{K6=|Wyr_< zKZ}^1YizOTSWWT!28qub7J2waMGW+1lg%%*Z5+n4+{+f{TD|O3{JvSxbqK%HzMB+z zuCuxqdHP1TBG25rX_2Q*x)yo*b(fIeIM^rOIpl3FB;P6IJB2*ubqx8=QGVl)?-KIZ z?-24f2a<0e^0wySf4h+H9`g9xHss^Buw0G(2x)V{SjV>zj5gjN-?P)NOn4nO+{#{QuA^I6?+eFvaTr1Mp zgbe_M_w5p9i1AlD0hj9 MX$N5wU%HC9fbqAZ(7|>R zj2L!y*e(f+>vPwHneX)LaSq!}FnrlvFvjiU&gV_@lV`73JItr?!Y1WS5R7Z9iGtD3 z_7Ds|*U3AU?k{@^9dp)yqs`m0ZXY`oQn!Bgv&=F(6w)MNllUeJ#ya}KNK=<>qT5U8 zU{ewn_lc9rIGtzIH}g zhs`f`sb(4d1FPrU~`tR37ah#^)g2=`sZB1usKgKVt?vH3t~|W zal6kKI@akc2bOh~qr82DF0&It+E*~voc#pD=Kc;lK(H(?tx`EqFy^y^1Y_JDEEv8V zA{g~~s9?!y{L*^iLU6gt|-+e4ovbj0wtxF+E`@ob@^&dzbzxq=aUz1(Ac`>-)t*$NX|Ye=|_l5&I>HE}rKv zO_*7!u^%B_R+wy3(PwTJY?GqATLeRQt6(*S?l!^j<#xg7Pj@8DmMH4_ z&V*S}#D14x_<6Tr%rW-}MmyXq82k471jCp61!Jy%Krrg%LBY^16l`2k*AEHCxK@5L zt{)aUVt7O_t~t}sw5;uDzek0Rdy9`HEGzuCH~ZYu=86XPRwTye<3ZP$gW;XeI5}kb zrO6b$Cdo7J*x~duPK%kik3YlCSTtu@U!XsvpAYs4jJTgjVjCV9HqS_Gj&y7W^go+V zCN_<+#ywRq=D~a`XD5X8bYStkXZd!u=zlzp(GJoF)3!CnIC@6-!*$oQg3))MD;RBP z_leOxc0~EU?hE1Ej)7C>&xd@ckZ0|?Kb+q=%HJ3AT|%CExi{pyM)`X}KEAJecgS}S z9N+E=`M563H`0qDCffTYhrKKqYt$qGxzs7w6*lo5i{K z^OiUl_D*pwZS{6>ZoS9{?7o|@_K|<@6|9layHt3tuG6WW@fZ!@d%X*uEAFo8JhA?%RT?@$_q8 z78CLQ67unQ`Z?s|@$^&3$K&b8kdMdH4$?%T2K#A^n!H zxZQp)nCc^cP(h9w27u@gdSA{hI&r37PbTv{;N zWf{S^pS7%D^t0s>W^GU6bJKVZL$|!6TR|{%D<;g!swoO=TCgTX-rksBR|;&S==fS8 zZ@$9U4EZ4uUnAs)C;1Uke)YgdM)}o3o;J6`tA;%FPM$Gi`BegE%vgTqkhi|Bd~X(5 z?0c(_k9}_$^0DvhhJ5V1^YOkd%xcQ{qFtu`=9w`&a>BE^UO0d zkG$`ebK+xf1opA73hZMi_OVxheeA0M$Jn)R+Jf`3uLI23#dieuv3CIWu@n2)+XMU9 z+X4I7*9P{n6Z_cL0`{@52^?eBn&<@1$KDIr*S#mOkG%)5kDb`Z-W}M--VNBt-WAx# zPV8gv0_|!Q z;268kqK(1%*f#}c?BX{8_OTBJ_OTQD*ardo*argp*aran*ol4Y{egY#{eWZax<4Ea z&c{9yn6Zl=0qkQR2JB-e_OTBI_OTBE_OWjP>|-bPv2PCSW8Vxo#y*ozLAC_vW8WH> zv5Vgd*vCE|*vC%nV;=|XV;>9bV;=+TV<+~pj|TR!j{;`wdL72tbdGETu8WIxJK(Os zs{6LqSNHV1@6tY>0FL)v^0mF$>bk{yF4;~De!79vy4(RA@9$(gDcJJbKNU)u0*|@yr*vGObu#bh<$1)Ar$Fc`7W6@j@(ckmWov`8bgwtHjNVqCH zLmpd%Kb!ij@+rtEEAnT~%`EY`Z^BKWE6A*bQw?r!LyLU8{gQrq?`8jltI`$ZfD-P& z67Hac(|8A$aEByZwTPsl;m{Ipp5UhDb8b)ej4w%gj2uc6Heou5V(#U-sHr<@#(5@CncQrAW zYuy^=Sz_+}S}fa?E552^{Vqy;PBd=+;(%`&i*mm>>8JOeFH62xpnH!#1DnP|ROJsy zHSt+382rtl;v&^F7IBv({S-^b;v&|xpK`b~=_j9;1#Xw8k9TLuzP-J~=N$>BH9u>7$#_XW z&Cy*6r`+!joWJJW6F818pZAt<_a&U>==I`8xoNEx%l%0|jrTypsoxe0GqT+0gC+eQ zD&Zb3;T}mi#W}l#d$fdmtb}_!;gtImfh&$-mILynaV=}G_^H=3PZwPM-u|iJhsJun z^-RL?9-i`S!c__epS>lV?wg-aIK38pA>rinMdPxZUM{$LPA`@C`CY>GvAADJIGw$( zCYfdJXvJq@UL0UlLBA^S_&L>i5@#(>+?o={@%Ml75Zk^u3kF6Pv!DK|Ip+TNQYe!xdniU3_{KV8tt}*AK)igI8R{D*^Lb z<73l3@S7pFwPDNpONf1JYXSS%)&%ykwFUOE5&PKM0Q=b10QT2G|5?^5=QibtG}b4{ zT!{Ne!s%N6X~L=BX9=hE`+34?&wr6{s=+@KPTxWLSHkVfBq3iWoc8Be2`8UlC!AvW z#yEc+QV!pO<9naN-zBy>1AZvD`umRGo2`GpOtwE7r}tNxn`-)#`O!Z4#jw`q&(_z+ ztG*2v&{O@r&K3o3`v3VUhE`^)+*+Eg;$_~dZR@~s42?NY!c}Mhnb$ZUYo%Zs8s-DX zF(mT`KV3wiEnD*ognsO%*cJ?&-$x4tj{UUe7Y>~7w@BdFPc`ND6#TyR+@dA@7AxT% zID1Z?nzP>|{j_ctPdKfiB@#}#w@c0!jkl!r^Ybb0x8VFL*wSie* z&1Y}JwfyiIm|@LZPs3W%>lxO4Ne{!iFX;}mxU0imfx7`~ye`1q9qtV5zk}5Y zxF>jxw+=8rr&IhLfqMfh-VVThfvK#|2EZHR5mmS!9#Mt+;}KPO03K0=`5K!lJP41d z!h`XMD!d6EQHA+7YVlU8#lv4Mri$MjkEp_1;1N}L2p&;|hvE@cco-g0g}JV4YrB^J z2w1-6!x2VFr zf29iZx<(b|eKA#d79LTBc{bF}huV1Y-lR4*)-3i?^xLWy1F(-{ z60o0(iNJm?i2Ynl0QR-n9(ZnZu~Wj;&M=(Kf%9{*OW^!?k@IsAIX@SX^EHT^p9^t* zE|kNbh||x-Uci1XrUCo8Aog>y2e6N0cVIsky8-*TAog=H71-BiSKv4owQDH)C0y-# z0Uv?$b1^+|e!R%}xrm&fi^%yJM9$BJI6oK4;Q++x=i)$MKNtH0`?(v4}0&NDwZhlDZw_>uE-6FEOO zk@Gc)oSz$Uer}XgoSUN&x1XD%fc@MM`?)z1*vD}Mu%Da5f&JVN`?)y`*w^Mz;JMAs zF~<41k#EgS?Y<1Zw;%lYxj8Oye*DP!xrv;go5=Z^M9$BRI6pVaDbCHQh}+N2DZqYi zi2dB04D91L3E0ogiNG3*RVVgya{{oh&GEo-Zgd}ZTHt)WPY;~;8986?$oYCl&gU69 zUvF`K4QcP6S;Sed$r%>6pBve7F8rBx0kEH&^MU=`5c|0~57@_iF0h}QbAbKa5c|0~ z8`#(JEZ{gdnu7}ir{~$@qQH5dk@IsCIX^d%^La+j&y6@gH@Ze%T*O(w@2^_ier{xY z2{=DDych6ua}}_k8)82$I> z8gM?3n}K~EHv#)RhiT43hkW7U4R1)PudZeSnl zUBEt8Vjt_Bz&_SHfPJjD1N&HseXO?u`&e%Uj@Z&ANwRYAL~=VKGr9IeXPVj*2jT;td9ZvSRV!Uu@d`O zX9N3K9|4ZB>OSv9a6ZpeXP#|`&gd?_OU(->|-VNu|5OrV|^Mp#wvd= zgY&We0ocd-I@Z&-xbe+zr)tY`Zr)7E3uFDJzyW}Ux9tB z?*jW+iG8ep0rs)}892tO*X#cP=hwiez&_SbfPJjQKGu(ceXJh=`&d5&_OTNCSU&*v zvAz!+W7V0pH#i^bI(Y5jW9jHwgdLDt_|#CCHAqd1?*#86PU63_i_}g zuFc~7{l;gOqd4y~a{hiJa{hiJaz4+<`TLE?F(cLF^AM-cVg11Q93tm)h@8(Maz2O1 z`5Yq0947Dwfq<7rJ~|CevI z=!L(DlJrv!UnZP#czb<1XqtoWN3KqC(0E@ZKE-Xj+?+lY%NuRxd%tgzev0$kgsXj?gU|Mo92CoUNk5$l*Ol};EvcE-#P`AHAf{yU zL*V$da(KIBo&A{ftD4WBO1PgBPJ8#460YGwI&ZqplwXU4Q!KBXFlRq;Et7szX%1^gto_*|7gAPXd% z?ss2J_EGJ#vWAA6OV;3miBI+0AoS~43#0hCP||O*@q>HuLDRZu9u`jeRjuD538!^E zE5yPsx_8<-skO#iH1VnN7AxU?Q^GA?!Yz?-x|W@m#`8kS08(fPS_;09VBHp>K#&sF&YvTO;r zT*B%7fvN4-yJ^j;d?14r&mol-A^ z<-TUY={p$>4QnNSbf&DGaLTt`!fE{W30JXLIwYLdNXLZhWBt}iIIX)*38(sYPB_Kc zCE?Vst8v|I{Jv?tSfl;z7Vv3Ni2Rb&pTH z2{4sHyeV2Lg?KZxR0{FtXsHz9EznXa#6!?hDa1q3QYpm4&{8SH!_iVH#C*M|*Nns? zfvFVYQD~_Y;?ZcS6yh;xsTAU|XsHz9acHR&;_+yy6yhz>QYpk-Lt00~6M?A|Vy=Jf z0b*W{bWI}WzS7-mmGNjhNRt zUF(Rs&$QQw`B@2-Ld>m@PQIh8`pJ+FOF%zG-`R}u4? zqU#DVuamk~67zmZ_e{imJxHYx$Fsk}ry$)+xE=|ooO&i)pPI{tUJ0k)?dqLyS|@!H zPIHp|ev{^3zl+q;#%I%LkLNkJ>EE=>epg8AcHLxr&DnYhr}d-pTe$J%bN!^B=4k`t zeB9a-8-nZP{B4x@(SGThaH`G5f$Qvi_ABA~2aet3b3no=mVtru;gtK-gi|cLm2kT!oW|QD;k19JC7kx}o(ZRU*vq(V z|L$FI_5Hg~;zxUHdcrB+8716I2@(1-6&Ks^_)$uBDJ)N&B%~to`+DxQ^xbnqbTIuY36G!S!;s*9BWpzs7!J!PRrS!E6;DZM9c#O8khs8Jyn(wKAh&_*fVM&aAzlYj-{RyWr9w^}+G|uOxeh(L1eGNTSVmmw7K4P{$R{4FjgnKOE`morL z#}iI-FhZL_ncac`FA{Ji;Moq32%q^?9Y^1By({mQKh0nEeFOD7AoSB7VHD4Y z@pNs>IKB4IxWRnVB>EgB0&V2_mvE$1moXs+a?0a{kOMi@<;B3M4{crwoXY%p^EHc{uUX`L%_8S( z7CB$D$oc&EJ>9s+etn%$!u{Uz^Ji~`t{|@!T>bv>^-v4iXx;rGaK46r44kiFOXPfABIoN8 zIiFuv7rn3XPRN1%l;fWhPG?yB9VtKF!C?)`A8q~;eEOP2&etq*zGjj0HH)0DS>za> z&X7|>4&=06-VHgB(`%W(CY-LLu`cXa;SWgdd~5n%fcnMx^Ywf$jOXhaIbYAn`Fcjq z*E4dyo{?i-+DF;BRG|UnZy^Vt-`@l0^NXC%FLFM=$oc#t=ktpkW7iyI`Dwkp7}k={ z@BNbeBIomqoX;$E>Igr!qWBvY~?wfsnAC%-5IiFwTe14Jh`9;p> z7dhsj^Nhbc;;)-NhYw3~h@8(Maz2O1`5YqWbBLVJ;iC{IIbHww`zCYC@8i&KADWs( z&gU07pI_vBev$L}MUG)8o=-xYiw3`?Ms7$oU*1=W~dh&mnR?hsZGp-Jg}s!)GNqM9$|BIiEx1d=8QGIYf>* z=zb}lgMKZ2UXnxPd=8QGIYiFq5ILViEW7^c z`x*aC`c>H(@~?!`drDs>oMO3lW%h1*zoZ<#O8V(Oq`{p7RZqPh2Lk#IUgS|*&vYn5zqWA3g>7YCL~7&0YB&FlbIrwO$~J zMfGfv#G-x+hJJoOEEG77E}si0oc8TC#fOzmbI|LtugA`rgM4l@LNbWSXha9T@CCY*BrZNgQG z4DkK>gi|a_C!AtgCgHSimrXe3zMOH3+W4zm%VrI58k=}^;7NR7vKp|z1~&q0-P2@M zU|lbXE5M40coksnXX2HC^_@@Rm4H=G;uV2aL*f;Hl?UR^5TSz`t-e;I0mL1?~o{@wx!>b6(-j!2CQ}{yPEl^JL+5fcaUk;_nE|&y*E! z2jDmls>^!Djid==RZ~9hhHK$!XsHxpu2-#D;x&M&6yi2$sT5)!EyeL-n3s{_QewCr zS}KM1{QOv-Arp51rc#JIqNP&SN9^mMrBcXqUa1sf&TDbh8s)l|sBJS}KKjGqh9+@#bi$6yhz=QYpkk&{8SH zL(x(x#KX{1Da6CkQYpkE&{8SHT&G&A#G`9(yV8Us=Ztln^K&*naLi9<^p+*uRtcxIwY71+zM8Xb3$A`nZj;!mhTEB~;@Lj2 z6*tjrb^cE<+j*IgiTqA7TV1<$G_3e{Fs!vZ#jt8U8JIQHSnr2(tE&b?8~40lhdUL1 z>*t-{gRvi7s)jp(f~)V{J%b-#r@a$SHdxO!ekC$`G#ShH0;$0fGnPB2^Lb-dYTXVQsgt260T!;1eT!`hRl7}gp( z8Q7mm4dL%l`7^2PcMeV~{MOGTf3Feyb<^aajhw&li5$CWUz{GsBd46tFwW1JxHG}E zr?JUdi66~F_T6v2PB}XnL)fv5jM1m5iyqcy;1O zHMq9m>TC3x65H#*`E_``+4{3f=jRQHAI3%2(sw#w~+VEbTVtF;>aK9tyM z{lz*woY?C8cqGIyJF!*G9u2mqC$B%Xb{{i8e!jI|9#1%p@k9xy_2zx5mQN=AG@ib% zx2_Iiv>NZJqF?!~@V$DZ~TOQYplP&{8SHgV9nc#G9a{Qi%O~l*Il$O5)97N2L(^_b7>n z;CU*Acqm#bg?JcRDusABS}KKj1X?PEcqCdXg?JQNDusA7S}KKj3|cCMnD0+gDa7N@ zQYpma(NZbITcV{>h$o<>Qivy_rBaCfdyd4s@7298G1rpT6EXL!_AN2@n)Vwp&sv?g z#5})rh7tSs9EtsVj>J=8Po)t1_Z*42cC~(qxpuXFiKhWmDa3oCrBaBwZ?t!axz@Gb ziMiIb-idkc=JfG(T@7oZ@^l;neS~ zgj2siC7jm#+a=sP30L7HLjIg^eQH$0UlLCH=iP);4u37--b*;Gm%ka;(#B__{m}Nj z+LKzZwcm+rx+XRL-;*&kpYJD}_Sgprr}}=Y*SW`TSDCHC#-irstgcwMe+C`D~eR>enjabY8ViIGqRcB%IdWLdA_* z(>N8&yh%TC^Cg_dn?K>yZ-InU4htrn`Yn`js>{L&r+$khocb+V!Y!6?TDw!WZ-6wd znbym1l71R*@q|;HFL#;Kr?@wdnUmu^H)V;$r{Y{P;i^S1d`B;E9qoBGO9hV4DZix? zPWxz?gwvW>HsQ2ymrFQ}w|v6s-eQG>Q(aa}I9&%;N;vgfIpLJwDkWScaJt`N?N&8z z9s-DR4}En9_PKWh_VeNAa$RtGoky!a!IsbIb@O@! zSAX5Oez3ia#wHtpyFvi%hGwgp9_F}>j9b*6>x@=0ZUK&1*Wt|#>pa=au-3v5!&;A< z8rGiQ1h@hK`oZf?ARz}+1l0Nevu{`&*>bhsaIFJQ&LF>r5S#oHG+p3{nD zXu>tp95O87v?hipT&3uR_iT*Qp5fD4vm?Q^H#;_?fjc>VR7u}4q3^iRcWlyEH5niJ zZe@MNZ<+K}%v*=P+ge}o+m!U(F7%xc`fi`})w-G(`tD$T#ZOB5Y7b8ieRs6J;-@5i zwI_E9eRr|G;&%>xIe)6-uHd-VWxE?N*Sf}?YJGJka$T#Adw}D*mapB-R{NK0TDH@I zAMPL3WiN2tKeF93*m7^{eX)JOaqq}>?_dk+*L^e{oL{>$%vSNyRINll|noOEtNt%6fKoPJPa+BLOdKTl|noMEtNvdb*i;WJPMdfAs&sEN+BMDmP#QW zim{V#6%KK+5f zyw}kE#sR?G!`jE}&-(z~3$Q=$*LBZ+C@{}9oo}?~S)%jg5b*wedh)zS(|sB-&mt;? znD=bDZzJY)T-Wmb;E(5=&bnE^(}1ZI+V6>$N+ISon@Twt_IsnHQpoRvmP#R>j+RO} z2=+73QYqwTqNP%Z_eD#k5YIwOr4aMJQ1^z!|8IOZxMktJY3&b3tUk8GfO+53&20TX zKLQ;0IF&-|_c^iO=fr-W6Z?Hm?Dsjb-{-`BpA-9iPVDzNvES#!exDQjeNOE6IkDg8 z#D1R>`+ZLA_c<~5IF&-|_c^iO=fvFOR0^@*=fr-W6Z?Hm?Ds7(_bin{JP9q8LhSb~ zvER4Ee%})NeM{{3EwSIX#D3or`+ZC7_bsvCx5R$m68n8i?Ds9P-?zkm-xB+MOUyk> zr4ajlOYHY8vELiSe%})NeM{{3EpdDep)>NxglnWZ72^GBj2ayvf-Gd zU&Xj%6HaH*aS5k>#|N&L`C}8mhuXD->si8)Qtl^&@jQ28!l~9LC7jm4$tBz=CETeA zr}O@_gwuGZC!DTRXBgMY)k8U)S#b4huE-8)>8aH`*>38(s9mT5?bWOi1;dFm+b;8N#HO8q@y_!ygT$`7p;eT5>EZ@O*rj^`;5!x|Nerjud4?P>nwTD`etkJ zp!Jmb{Q~I!j(mIQe`j;gtK!38&nDmvGAcm4s97uO^&w z|9!$K&esx7xxb!p%KZ-sC!c>bPHTkALu>7gf~%jSODErhlI@%3CtJhI75(aG(Vxt( z&e6A$7*xl%%~s>Qlh}&;i`lBqe>Pj~Ie(63-?dl0-VOa&6UC$NkISd$_`T|QAFH+U z*TkpJ=J!gt%%|4o-;#df{+@8k;r$Yh-zc>3{u=aH3Q)}(BBu>qfes939HTkT4 z-@obS)W@wJyVpV|2TEc(Uu zm5(n!Sxx!<_PF*{w|^zE>3y{?6HYmORl##8P;2Cj>(Gd4d3j?XE-p99yKq)A!7Y)xq1e+fS6PFzEk zC|_qSjPvVNb!k;__4hejCbr6{wb^RSc@kT3^O>#o%)Dmn*Sh@9pK$88Knb^C;QZVy z6u1V6##^|ATO@Ej*wLiS=b}{_)$8y3#rgJYq0HxE!6!#h4!=n_ot=v(ocb-1aO$^Y z!l~bH6HaSrsScv})48%-!s$F(zJyyL;dJJ#m~isBQo^Y&`fjS< zTa|{I=q{OZ8m3w~Uk(QANXU*R^u(N;0F1?OX2(`Tf0h&rkwJ~UjM7M+4r2ama_Sh&+NOdx^~C^kARVBymnQhrg18V_6etc z*?5|_jMID19TK0j83LqZl|P%tQapuc*C`(4T&EJQbHYu7KT4M>f95{BCY<8zR>E~p zIL*-qCt!3n20H%U0H-Axltxo;LY z|DN0Cfz#e=fNW914GA1SJCqEq@=<+m6(gS=R^`u}8y+}XY92-y*OE_~$nVH%{re51 zMgh-ljM0Jf`Hl&k=7iqH22Nv>8y7fVi}8W;{k9C8?tS61VXMHc&L>T_uJS>>hWuPs zvQ3qb>Ku{Q!M0VRI_Kluu7uk@aK2^}0;l=HKn)WEr+q_iQs8v%k=r3~ejX+VPPw65 z!<4}Jeme$^_Nwr142z^PTD`n}OLgeJRH z@i~qF+TDS@Es@TEJ*rLbjiv$5&34aVyH{eX`P$oTh4)Ep6~px4cSecr%wW53Vyk(W zWwyfmCAPYU+23r14@hihHq_)mvlTwbY+D)M39a%z7+e=%*&Yhq6nES*8>jiCrML^gwdWHi7gqV8sUOYXjEm>wbjDp|{rni>E(YhvsG6Ad z4Ir0R`Lljb5sAC3+H{P|fny%3`4tJbFU=uWR{68(7~-x1=kvJQY_mMBG0xXRYvkI3 ztH(aw>YCZcF|R8b^ZI1Wy&7twF{c~8p~N&MhNcYESTb9+a@)z7OtOKk4~=f}L;Z2g!s8)|Zoaej`phwcUE{oI$>swVdb z+XoU`*#eJHWj`hPgsK9bmK9nTK7k0!Rt?Xh6{cw(z{9sNF$*lPX7Iy{-!>il>L zoUiZGiLGk(Ot9T5IkRM|wc+Pmc|U7ocusrixrA$MsLAujjWoO}S}KL#zg-P2l|sBa zS}KKj4YX7WaT~N$3UOPsR0{FhXsHz9c4(;-V%ANyBkll9r4V;SOQo!TQS;v(?~giA zDXh-`v{VZ5K(tf}@gTHR3h`jHR0{DXXsHz9P0>;*#G9d|QiwN4OQjHRftE@k9)gxi zAs&jBN+BMGmP#QWj+ROx9)XrhA?7;OS|uI@Or;QyMoXm-k3mbN5RXMmr4WxpOQjHx zM@yv;Z;6&lA>I}(l|sB7S}KKjd$d#v@dUI~3h_j=R0{DVv{VZ54rr+q;>l>K6yhmp zsTATJ(NZbIyyj3T#5~8T6yjabQYpl{qNP%Zr=q1&h<8Ivr4a9qmP#S!b(KmXo`#l6 zA>I=$l|sxjN$1iF)!KdL-XtO9dGAQ25cB-i8BEOUhpr*Syx!@WN6c%6t{cR&fT;X;tkPKDa3uzQYpk6qoq=a`=O;$h?~#O|8vgH7pwoX&Q7jh zo}0~QXKgRHK`p2h)}}35DusA$v{VXlJG4{^aeK5>3ULRtR0?rNv{Xv-eP{F8$$L&J zh3jK9S}KKj3|cCMcr02fg?JoVDusAFS}KLO`M&eN@4oY;YV+B-MzPXrYpwb0Y`*W@ z!mWoPcub}6ENXtu{{Q?nJI@-vUZ+x6&;R7t?2NzpzOz1;&1Yxx+1Y${*7pP7>!eb6 zR{T$Xzw_nlzxRG;^V!+_SyA)*oyE~moDDoTo6pXAJ^qu}>v={lEG-`@eg3@@$`hrADQ&|4g)03h};ZsTAT_XsHz9=Fiz{wP^mVsQI&^ z=Ff`$|ND*-&$9oO?K0E8RX#SkN`8!I@-|wvNhyUd7cmA&W zf9Ct0Eo#$O{H)mHSE_u_^!JV`dXsHz9 z4bf65#C_3HDa8JJqQw4tqQw4tqQnExk4hmPh?Ytr9)y-kA@<*~C-&d5C*BlxR0{EC zXsHz9&CyaR#QwEEvH!fD*neJ6%=4E@As&vFN+BMBmP#QWiIz$s_V=&EqwzeILOcd7 zl|nogEtNt%4lR{JJRU8TLhQfmMC`v~PrM!Ms1)Mu(NZbI6VOsA#QrlU;z@X(N+I@V zF!5wOPo)q~K})3&`_Jo%{deq%{deq%cR@cYg?LxAR0^?w-;{VaJWr(%?~ayAA@=vL z#MAIRl|sBHS}KKjFSJw&@!n{u6ykl*QYpmK(NZbIGtg2g#52)SDa8AtrBaAzp`}uY z{cC?>{~dc`{~dec1JRF4A@<*~CywW=&eGSaHEuCBkWNsJP5E23>F?eC)v)IEUBjyL zUktO>l=rH&{;J!b4XbYN0C$0%@Y}#$9sUz=H(-tT7I1fm-vsUfEdOr+_jLG=z`cOA zF8=`B8(8ta4jk*EIsbdYRcH=*zsjFYuUo3;2MO1Q=8zAo{F(RrDB%?6$0gh+CEP!Z zYt>K})#B5FtN&ihXNI)~KQHhh2WNf12z~z<`u@xMwrr@2et+i6VEdKXwy2x_I$vKK zm&Ndn`H|mm4J-cdtgrX0b@qM2sV+YRj_XVEW0jBU^XIvr0>^bLiTykm`>~H=`MF9| zACE}fFIA#C$9*A*<9V)ZyoO8Yq3)ANdn^7OZ9Pw$>~F5?_orGUK9xhuz;&iGlld>; zgL*7HCnS;M^YYm$`1D-s!1?*heCjNiC-n0^=S?_X>nQUjoc8qm38x$uNI0$Y1ru&Y zZMcSo5^i7P7B>^#S2REJyOLqWzq0jh zA*X0{hOLrt8nY5Ozs^@J;Ti+y&#BnYbFm-8RV=FopPpMiaQ>W%<9V)ZyfuPPj;FIK z{vAC%Pn#!_dnBDhhI_eeOM-#v}<=aKU3RdDt5d!Lf?yLU2%;_g$zWq+?*XVJPzKdqmvpSbmseyVTT z-;}@oBF^9;2e)QpGSS+N3ZM1g-`>!$UK_4&Sg~)A#Gv!HZ^Bhx zjn^;fS24TV-|TLB-PIgte~(%7_GB`i>Ztnr>rbrJzf&8gYlj(ojNxuPW2 z>~CW$#{r2?aq(|sa~#!sU`f9$2d&MF(|CgtpA`-Y8JuwHw@Je3tl2c-G>4ldoP2Jc zaC-f{MZ#(Rh9sQkZ)n2l92%Bz(`%_W3{N=u9FcH}Wn|#^9)x6}Dj(JNgy;14to_=q z&>S)fK0BbT5I~66L0ch!5PQEws_E#2(_D-x;l>88H=UZqJUpj7T9|D&bwV5a`TL;A z`F_!-?-w}^r~NW6oO*2P|l zA8~t|AK`rxKgxM};z!&J^CLXdY_*Pg->SS02G<2xb8--HSBDP-?gp&a!v_F!Zwc=Y z+yhwadOu+Qp5&DC>+2R~4YR=WEKrX70>^Vn_Z8XtP`~VNXlosRl3a5u90qboGG4V{ z8X68wxc!V9hq>7 z<*0;{&!ZDgYyFsn)7m;V;U-xQ$0Zy;W1;jO#-B}Vt!w9sNv*3iha8{yRDLIva3?05 z@;fQvw9ZaWIL-Sh38x%RO*rkp(-N*yWY^Gedcw))840Hx&P+Jn1D};}6KuS*6Hd9G zlW@x6+=SEla9#;_ehGI$33p)$cTvJku>3AgINjS-6HYl?l5o?^=cOgwWeKPA?D7)s ziW2V1gwxz!Rl;4JaC$9pO~UECzc%5tCaz03?dj_iPG`XlCESfA+)X9i%?YRZyCvZi z%dH7F!RGe167Ked)Ajd`gwx)-w>RtAB?5XyQk(eIoHA+x@D{r0L)D)to+Mezb0$PW;IBnS|5)KbvqG|G9*d z&*u|PbNoWWX|7&OIK}c(!f9@@ewvFcPOa0I6Q9cAcO~2_30L8CKwdS@uQ%1~_u%*% zM)z-(LKISOCDsD>W} zKfG;M-k*TuYbDu!9BlbIOxK7{!SVIIZ2w`lE%2N~>;JQa)9d@sja$?@aSYYuCve2t zD?b{R-yaMs{+|u&y7|3f-RFG=%o+=S3(Oh|e*?^VXuPk1Sr6f_fLV9>{}PyW7ycJ8 zYohr73G81xegPb>1G2rcI(P1>gi{TENqlO}G+a8TPjQ(~oh>bre&SjtoN{PY!rhhR zpmX)^60WuR^q(bEm=Aj1%s}fGKs{9vC;?bF+`OV^4?0+F1`ORuRyUOgE z)?9V{O%kil@_sT5*<=0v3s zcSK92H2)rJ)87MNeFmT|R0{Dxv{VZ5Ahc8p@nE!63h^drsTAT((NZbIo1vvrh&M+| zr4Vm{mP#QWf|g1l9*UMqAs&X7N+BMOmP#QWftE@k<~r3{B_0J#r4WxsOQjHxK})3& zk3~zR5RXGkr4WxtOQjHRiIz$s-WDyDLcASXDusA^v{VZ51hiBN@kF##3h^YgR0=WQ z&!SR@C!?iOh^L^XQiyj%OQpP2ZTfov=_3h{1esTAVf z(NZbIJm+-Q5l;iAQi%6NOQjI=o`Omt=D9_s5c3?QQi!LcrBaA_)>A3OGtp8h#JrwS zDa5nTQYpkd*LAiN?+;9+5Fdb+N+CWFEtNui5LzmQxcTg44QjRM##LZ{t z|G8%;=b!hVx(6lxPd+dN+IrumP%>94{bh6xsJ7#xjse%Qz^t_ z&{8SHW6@G6#N*IXDa7N^QYpm$N%x_Q_kYq^%6z6{9;g)J8EB~#;+bfv6ykl+QYpl< z&{8SH&G(_TS~Q=X&1Yxx*;#)L&$Hpb^Xz=7+WZ>6`89m)?B!bPiicDRaW}RFs^RW< zOr?kO5yt8d!JMa zaoPKyHSZh<^Ma}ceW^Ir)ex}U&<%zpNZ}RM5VL~+(vk=0;RMM9G|Q3 z2c&~>t!l9NsrR-!g464D5}klMp%uSQN#D+)Z`aVbOVan(H5vN$u)gBECw&!D&(OEG z^%dW%q;H?lcfHVe-K4MHr(Zwx-O&1q-yrF$_r*2}eK)qg;`=6j^}48E=sUpritiu# za{g4uf#A9T%XTm@uUi^-ko9ee=P0~Jsg9e1>j^Aho0zTEEw4|q-7NUwwNQ200vxY} zvfVt`b_S=!tgUJ?2AFfGct->Kvxlg{ACO7LRSH!9ynt|TaDCBEFz#^R{`Q#7UWRpj zIt+NA+`omfVXh`9)|UK+rhv?9KX9^<#`bBP{;3PSo3rs@JPo` zHLN-w0IXbTwyR-1w?D9QCBKVdJ+~jQawWgBVLdksSh$&N`%9Z>MhV|S&z;T@@x5i86&K(i@aqmbz$M2&>pIWC!hJLzF zW{)4M{8_(VdhVyd(NaE-3O;qOO`GV`bI~XFt0eaGTvLwpX^gWHPBlEcggd8%J2&Af7Rz}Fr}@+xZeg};#GN1d z`Tg_^J_qw>0+I5&AoS}+(qz?ec6ly4*W~lU;M1?I=+kr2r#~-ZKhMQ}KKEx!a=$3V z>CcNep64`Pc9vXR_^h8LRl_<OFn6$GjFeO=G~b1(V2Hs!fA|~6HYa}rG&e+ zgu5-_8ZG|Yjr04dV%&$t{;hxao1f`Zw4eHhGm4yI)!*i6$tO+pI_eIL>0j61X)&~x zAzIDH-33?wzS&*YSNuJORoi>bS4-z>x;?(nY=!SnY!%}JW-I)l+4}RpVtyY=xC4!Q zv}o)1goLV}YUF)t9b|tmNoPgb-{JWnIdj!dKIgO8*l4_miyZ2&1s*r7zK?It;y#?KkxJTgj0=QFm7FTFwuNHW#^4vr#&CepQjDW*E5DSU(XrVd_8Me^YsF7 zUyMnmFu#q_QYpm!&{8SH{n1h>!~@V$DZ~TOQYplP&{8SHgV9nc#G9a{Qi%CFoJt|y z3@w#Hyg6Dbg_y6)sTATNXsHz9p=hZT;$dj16yo7%sTAT7XsHz9k!YzD;!$X+6ynin zsTATdXsHz9v1q9j;&EuH6yot{sT5-W{So2`c%Dijo`{x8A)bVmN+IU`1eHR}wWRe# z%(bNTM7$#~l|sxtseMVzvr*?GG2au=`vSzf0#hl(yjJRZNzC^M^gaPG&l;UK#9X^t zzr@pksTATp(NZbI+;7@*#9ZrI@5Ef|TJOZX4(M7yJOh|YA)bkrN+I4CEtNvdb6jUR zvHv~@vA^%{N@%iBl@ID?VNZLWO&J%@$zCP>JQw@1o96As5U1x}N;u8W%f_{IF=%bP zQgHR!{VuW9x_i}Z70>SzTXC| zSb4u`STVe1SnKRhU4*Xbp55UKjh#nmUqpUV)?+ZV)?6K`F_u^ zV)?sa#qu}cb}r}lfqiTr7WVaZNz;nif0S?<=i?IYlZ5MIKL3$$^7(1P>FoHdg!|mM zmNo_(_4@)G^A-MQiJyNZoX(0bOSrF$Yw5;Uzpue@eBo~rKZ@W}RsDVf*WNwf`m*}9lTWt-*0oT0q_b}cew4#4fVFOmr-5}=3v2F)f5G@# zx5Pguc_{yR;LD$mIh|dS*$zZJqaIc8`%0-R|od768l(J1NN~t0{d83 z1@^HL`&cW$KGs!$=N9WhA=X~7^|AH<_OTNCSi1xJSi1rHSi1uISc!eCU4VV8oq^{T z>oFnLzOePNZV2pSCHApy0PJI3AK1sb9NZ9&VhXMOoiG8d?fqkq)fPJi6 z0Q*>peXN@U`&c&vjpeXOH_eXOH^ zW32M`d5D!~n2&WgU>_^7k98`rk9AjI`R1s*0Q*>peXKhJ`&f4ZW~}~xt&vYb7D&!; zy|!5};i^S1+%G4buCEIxoYveT38(QEO*oCWSixsn^ZnovHLmgwq~tpK$Wop@i$0aH{7z38z>(m2jOCPCmONoYrU8gwvYnmT-!td&23xwH^tl z_Yry~oP735IOWhg;k1_eB%FM%n{di+y%KKygws0PAmOyfHcU93SsNvs&howqr~STh z!l^FV-#L`een~&A&qprd#%Ox25;r^HwC4LKKE-AI><_V{oZjo@RTQ;87O~&bZJ;Z!w`Np~pO~!Au8V^f2)pU3X zHzMKmp7F>MZdAf4r_m+cd&wHod$MDae&WU&=j){R)y5ZG{k`9Di67Z!eza!fN9W6y zi67Z!ess;*+TxMzR*4_kW`8$PYib+wqqA_^#E)#ZOE}H{_6eu)CnTJFPE0t>@uY;) zTie+-bX>PK9nu{z>tCD>MIKSRhvt7XPeJ9~vOKhhG+uagd z#k#xM3h$BFYRqZD@17;Ldj;FQ6I-pZeau#PdSa_{ZHCzj&rED}53{e?3ePgzR(5`| zp2~Yaa9x09djN1(V2!iC_0`<;0H?Yf1g@v^b)ebm+S=RM9vu911E;(X1=ro#9ujOj zgVXioaB%!gKrtLXG#mcU00(d-2NZoZ~5$ z`;2Q%6O)P79#{CR&&RQb^;zK2h86oU1#X&`?w5~GxQdN)Lg3cH^Wshn9J|Qpy;Y*7 z<0+Ptl71DMLl#TwuQPn1gj21at+IF1@l<2YO)Hw1#5(>vbyGY`C-LZPJ~_!zYxVjn zgUDl@op8$Wl*FgF<&s!b?^8?qWjScwW}L=5&3v|GDv%1akq;pABNna`;rW1>Z!0z? z+NY<(PkTLy_Do=ZuAgD!WPQ&Hea{Jf&rbR({&Pd$^R2J=^OC-b`GU~*BI_&u!jirh zhrXADzSX3!&cjPX-^;D9_{);Mn#(If->a;z_$!mXI{U8&T5}t25nSM{W*&y1}37aw|B09l0gg`s>K;;QV#uwqVOJ zb*|h2&ad4&%~tWzR_p1mgsad1a(4-LPr~W_yL*jO?dVHua6R~ zQI7Y6^Ep0{#G^5X+dlXC>;CY;WDIc+g)ts38RLG4Z1syj*bg&p%!&vHcx5zYbq9Tff$Hmc5$z(cJ#N;Ob|?YbCa?gY#ql!EF7Q zIs^Zh_|YDEBlvkUu~kjp3bubrY?aH~!S;ZNXvY^^SD8XK{X zttGIJtp%`;tpV7_M(ktzCB*h~5}W3AFz*&*vB?6 zu&>)Zz_D(MZ6R$Nu`LDcV12KKQM`&d^3 z_OY%A>|K@6zj_3fcCv0^lv^QIw-Trk% z$Kaxi+quK3pxoz2#tpQ=?CaK7%nfPLMGecgKk`?~i4mM?_c z&>h&1yl40ADU0b3vA z=D`1z`DtN37?a`}g1&xTv;LZcp@Cz)Rr6tiV|`WQ;l?RX&y4`b^{F*FGWhYajVj?r zhcTGHYA`m8F(%j!fUWW$2hNYFdMI|TPt{{QY_s~vmg`gX*%G!{y=2=LwyM`wu=O=k z{aV9@;%&EvE$3eK+a~m#TYWXI&wX3;^|`CwneBGLR`qXPL|FTMCe?oeY<-N|1OGeo zQ-M#N#S_t&kA<_}{nfcX3ATPM6N%fQgqxgjS|3va$7`hI$0{GyuL+*}DR8{*Nxro2 zC#udB5omY9c*>Pe?-<7CbINJw!1*yY#or(EvGF=9iGPcss|d7_^EpM1-E@uICB)*n zT}!yBCERWar`&fpPBr9snzJ**?+>4y{I04#gE`l@%+LA3&xMJfMh*h4s@ldN#uM^k@G&|@9TS?k@G$y$1t=n_6_-YZdM7m zUkSH=!l|AI7?-W*`@?#EFj>#q7Y~Q|W!zd%vxA?<5A8CD@1fH)KmE(!PejgL5j%Su|XC<6`o^708=X#xcZo$=G2cMJJD%SJNR`Hyl z*owQ*Y*p6_%+{{~`MoINRPT!e$687B`@F671lo!{9#x!&_1B&JI~g4UzfQoh?uxG( zV&D*}?InTZ+K?zd<;-;=i5&N@ByxOS?@eDC#v`X#E=xG&aCyRMZmvi;z23hv;q)26 zRSCB*1A<(gaJmn_CgJ4sTI1%cxhZ}shwBQizV~i6to*Mxth{b8ti5-WVXdtjfmwIi z-%{}PYbZI@{no%SFs=35O1Rq-PVwGhT(&RoD!BT-yfd*?UG6qp#dA+$EABqCRnGUC zt=1^>s_+No{=l&w^7BC8XfN);67HddQ#~FwF004vf~(i#k;GQ@c+_ka&tqnrt*QTx z-{cdM}nWq28yW)Dn;?!I`p5#bUqo)P~6I9|t;``ablJI48X ztKU4AHLdrbVas}p|4ZUmF~3{F{k4R9FL1hFVRkdZ>vf-hjqjs;P zGlY0GJWr(%ua1^VA;!O7|8FqjHo#O0aa*)h3h|m~sTAV1&{8SHYqKp-4fF3}=x<|? zXKu=mnD=v33UNoYR0{DrXsHxpo>w}vh}Q+CQi#_>OQjI6kCsXy-T*C?LcAebDus9> zv{VZ5wrHso;_c8{G3i6T7yT4XU zs8J28Z@~xkdmbY7`#9+*?voPk9|@<|OTCH(l8^Uk(r+>Yf_#>6)q=ry2}-yx5>7e% zGvTJ0&##g=CmHvzq@UvaGU2r5zAE9qPB_K#O~UD(;oF4M+;dIvj zm~e{ar-ajb|2g4wo%^je|e3KqTTdQlE75>E5bGU4=oO{;{{c&&|VX>-k{(H_rR zaP`;n^CW&Io9%oFr)%Q;38%lwwLrot-vtwHrj55y!fE|1oN$U|k%ZG+FPd}VigUSyQ_Yr7IPJL= zO1KphZnEXRQo`+E+kX`Qca{DZX~VbF0yqlKPEJ zIK95;mG~5w^_$EPAib?$s~RkRR_t+~f~#Lg)-$ZSt!sT-xW3}nPq@iuw{4P#`ei<~ zhqHe2IWF<3e&Z8Pb0My!%b{YA^Zc8xZTV4+Gd~*?wTN|W>PNBeki@F%+J+@H-6-Ky z)4mBe*>cRr)4ItxoevwE&n(|qw|>d^jW&M&gi}oilyCzRPVadRD&YnvoO0TvgqxMj zi{3NdH0dX9Gvj=nbgpbsaP{AB*gWwg+su#FjQr?&JtXlX+su#d4~AJhvK^ZEk?nM= zhtF%W&FOISqxvYS+=VYtlRwdll#`$}aiuIf4(%J*vAMkft$!WjM?ROHhGhDxC>gQPX-v%*sf?wfn zlbCdlY+rEo_XD@HzTzhs*1DT$zWlhV=OnWg-XXD7Oq0!4c#7HjTv;2+jtQsx(VYTU zDSFg?#NP|%v#P!7=f|6DeebRELDOr1)>_%`o3<|IH}6;Gb7#xJ*Hq)}0#5HwF{)id zzI;wKyfZlmwVtOYV~E==jIn#k7<(j~_V2W03~_sgG4@LQOlC5Wy%SF7z&^(LT8Wzu z&ev*2;zzZbnfTF}v2VfE&wl<6CD(%=Q?_x;StVocmyD@>r7^X4_AjxGV;+!9+@Y42QJ!d;Va@|lgN zeUtfAj%B}Nnf-32)@bHad+*vL&dF9c{hoEU?rXmv-SqxL<7I0`XI}QZ(du_y5{uTa z#`D+3^7(BVoBU+4Dc|dp@wF#4pV{~~6#e4*%EyKzZcR9?-`f&Sd*b$ln{2UU`DtG6Ncw3{WInZiv%2V9%f{3BbZ0W2a=$C# zltVUuS|fKS{VJCGJ;u#vn2mhiTX6OF86PyP+T3TD~c9HsdK6wpM50slmCpm50tOXPgL#rb+qX6TUDia6`9m2~F#ddv28*!p_^3E0=0*w_0lU?1a~z`ovZ0Q-6q z`+ENo*w^h3z^u2AP5v0SkL|C(J~m<>+q=L%w!Z+&7f1Osu#b({$Mz1ekL_*X7@OAV zd*FO*9|8N=h<$7y0{hrL0QRxH5A0(j_ObmP*vIxa;27IvJ_Y$0oR94bU>_T?kL`0{ zAKPcZKDJMReQd-&wtoQo*ggS{vFRHAEjSKLPvLegyWh{Q&G^BlfX<5A0+64mig4tG_Fo#x^go zkB!*JHV?3mtu^qkzA#EFU>_T?kF6!JkF5o8jBN^^g2cKl3|k)?v5##bU?1Cpz&^GG zfPHMlKDPOReQfgq$Jlh8kNb29*!tLreQb*Z``CU1>|M?j7|6C z%Yhq0+GJT^9~-fcZ5d!6+tR>3wxxi5Y{Wjc-vayCmIP*O{&n_bi$&KAasIx8e^<)) z^FAZzuY-~E*TKm7JR|3?gW?voxczru+8{sTDaNf~ST$SSuwI+CHLRMgW?1<(0(XHg z;Z=dVI$Qzn=EhqExVytE1NQ)y|CNCKcVSip?gd_RumZ6EF3j@4zV7<|8=qY>aEg)K zT7lD9Om6MK`SsJzINf9Dy%hHA08VuzZXf)3u4CXPvV+Mwftw@(ZKuF#-_WFU;C$R& z0yoV0>>9XHj*E2~+T3xh(-4Paoje!ow6W_K>(tM2u}*^=7whEXj&csV=HPt`uRKt1y1K6 zx;MnXL%?%eGWa4ss$Zu(H}Jx_bDM;Iv{ao2hgitz+`Tw_r`E^0W7))F(Hwf8 zpH_+LbNG3csOe|e8vl!etK0s|Y*lae)LQt;xNOXC3a)PZo!NRnTAQ0#J@ohKI$;dS zX5f6yM*?>diFO1quY;0Jfx8Ky9S+Rvpkx?u4*|48fq5O23<2hKP_hMZZvnKM1N%Af z^ArD;h>!cqOPludT-!_L&RrGy(NgPZREUL~=JAP=bu_(X9j$G#WOKS?!PVz=yi9kl5;V((b|b;>1?pf!HJ1UYgkIyy=g<=i|M_Y+Gnj(dt?}4SwT(lb<~k zPUqHM38xzDop8EO*(c$2e=|Md^uEiCgwuYSnQ-#CZ^9{-S%LG{llZqc{Ix4`{<;)7 zR#)@AUl@;^>bZa5{B>vhaF+V(YUI4n1AA`oIh7p&m#)1{ym|?6I<2yNV8QuM*`|i|))4n?`;pFpl-P}kRQK}($H2s0kZ_9W!h};k7bToxxj5l8 zC)I>g{Vqv3&B>(+r#ZPS;pFpj{P6U19Md*MnQ2_OSSQ1F&9W@Co8Y0$&NZL%_2gzR|{RZO`>}+^yh-IeZK71cz@1 zp62jP<|{i3ZY$w#Pq;?&dq=`mjJq@8GzWK;aCet*_mpt=8t2bjjdve7Kd1L6wmL)q zFXGMv$g1KA_X|c;%sE~$E1Iq33JXl=bZI@ zeYZcixA}Ie{#CEuo2tvZ|y z+EXH)x+Q|+?&& zxR$;WjGwE&7L0ZIMljC%w}SEf`%W;<-}izs&L0G0U49gdbNiEEjPqx~I7hz-#x?k> zV7%A+O)%#0yI`F6KML$m!I*L_adsk8AgD(T~>;_M6Sus~y;Hv+BBT_&4f3 zUGqEARi~JB8#c;kIav#IrOum&~I)|b4wYx3Eu_KnBoXM)+? zmU<3h%NXZ$Tk5&owiyG&vEqN9TQIK47J_lz@fjP@C;QFg{aW}K_B{yaWf8FU&|ISn zL)-T<)C)n^Lhp;0wFRL&;z+#!bSG%8-CfntNE$hB~vRy&ga$Q!E?TW&ddvB{` zyOP_+^MQ3(IoYlPZEJzPR)e;+z&KWQU(J0&Y~ER?HNb4%>3emz<$AO^r|p`_kF77( zr8StXFWRn^Y;EoG+Qqd++cs%TnEtnp+X}`t$mf~F8gYMZU+vdWgLba@yqdLL^L$xH zjKh7bgJ2a460KG+_S>tvAL{4lFuwmK9)5SxQH;r&sOM^MZo5yhh#S5OGDY_Bq*+jK?WjkYY*Z!5ug{%tCJlBs^Iq1R?3C5cC7L0XM<8j?6#{FR%_o?#D>t~d)_|Kwdlh=Z*SG zi8wh=e1=IppXqaaVLY{FxJMjZ=y!mKh3l8&sr^;YZ{zC;{ixWO?}1``?t`39HU9W& zzr4Qm`1GS{YR}rpdjjh=LBz&=ZlYk!>7W8TSTL@&Lj>dcJybBB6Nd@LSX6$Tm%~Lr zo)gL^*RQGz_cb-%AYb=Ki1C>Fk%BP?HGf(kL-BSk8@OdW(oiAclEjO82d;I+ zINmwJXDhGixdnEfvv{509)3ZUHQc|SFZ|HVvRA!DWV_Q-=ZJq^-9r? zan62)O@8COmTX?Zcn-~&GHXB9b79esxla|bFo*nX0~>+kT_yUFO-qbpS^ujOlkNdG0x0loSDTqGmCL%X7j*xX1{A2&y|?Nb%k1I7IVlf=8##;A+wl6W)?NCVfmbn zIb2`JA+wl6W-*7%Vh)+b95Rb>-cX1$vlwS)G0x0loSDTqGqYNAKf5u-8FRQPv9Z?C zTV^qb%wi6i#T+t=Ib>#GFrJ%JoW_`6>3p~)^`k%QpIOW=vzT9IF~7`WewoF2xU~>x zW--ppVw{=9I5UfJW)|bTtq^BsG0x0loSDTqGmCL%7UR6V5NBpF&dg$*nZ-CWi*aUV zYmc#V?Y6Ktz45;14&>0AZ{m0#bf2XA27NDhjOiX|UZ2hEZfIUZsQEK1(|3W7wZC|~ z6WXpbDz^K2CNdk{Lt_5gH_jr+l)U@^9*pkr*N zV{A`C$Jm~LjcV|xrb#x?^w$HqCG2^M2}0XoKJI>z=qbd2pe=os6x&@ndC zF}7!*V{A`DHx=89U@^AWp<`^OV{ETM$JkzljcV|y7o#`Y3)j*a*7Z-B+v z-h+;@nU1l&3ms#72Rg>~Hgt^5bd2pS=os6Z&=y<#ck4WdKLv~bcIPwb80#m{F;>$t z){miMtRF$gSU-f0v6_ytegGX~eIGi<%KQ6o!7Ntl@1SF>-$2J$O~+WjhK{j*1s!Ak z5<13eI>!11bd2?LXp5Df!!0(R$KQjs53&9T-2s|)|H1q6b9g5(o@YOUbq>CMa$9~5 zwx1Qz_LtT?Y53pGGnyVVu-E^$` zzp#&W{|B1B%=h2WvF@g0-T#7)b^jCE>dxAkqJ1-r9c`OKn=NZ+iuIXIFvc~zVEmqJ z4#9?+W3)L7Y_0;ETQGjd-9j*4_vaC;e|;!5^A^~A&YG%WoUdF%t1oNZ3F{`-*J{i9 zwoK#Oxxw>e{=^1?*~{7jiP`u0{7!$t#G=m`Q>=HxT-*1;{BAfiTW?%*T~ZElyoJ(u zc7E_2SvWDP8SAo9^0^Hd>#|7dw<#FwvS?zPfUz!%CAK{n>#}%a(Pz$mXY^~TE`7jQ zm+mQtINlOzye-g=by+eo+h^Fe^k-Vcb;sO~pImRn*iNkKhWRwgb8smiqlG(;V=N8U z-ru)FALLl`mX{V>%X%5 z;vO7ruktxo>*2Dr9wv$Pz%{pO8pC2ETdlxWcQzZ73H{apYaQNS3)=Fg-!+9_en-?= zFs^|%1=hB})-JGh1Y<1i3aq`em_PgNi*+3H?*LoNpS)K1<=)Y;z&aIJXTf-_>mnHE zqHBS5E3kE)#r)~h)>O=Yedsv<>j}TSf3W8q)9@H%Tm>@0r<9^w}fzi)_=x;vTGgGWX3=zvy#w!C04` zg7NyjMPhM(SI;0|ep{w~77N$JR*A)HQ1%(GN12&VuBEM$&p2K$!5C+6!C04V1Y@nY z6^uXk-A*vB>pp^UUH27?KKnVF%M$a(9Ja5rhJS~*i)+?<2iMGNN7vkkcXrJ=-wE2z zZ`$u#<&Ez@XuF$WgFU|81!K&6Bo@!9_tNz(o>Q6GQ1rQH@)_A)iN$kD`Q$wJPyM3L z0fKSQ99Up`3&y$H$634&7;J;04XUz+&;9!fTh?N*+cKUZ!j^15x8-~fbz9!2*w_{O z18taKgX>Yv@B$kl7-JooSX>LA7uG^%76yHeNCG|HiW=rn@gY@L3XEzjcAxb1E!G*0gP4han$+Ua+?hO}R(gBa-dm z1>2**dWCV0EZ82CY>#$Z=3;JG?_&$>IA{Fa!q3ju|0FQ34bvw=^Ezt!1n5!r7jMVA zFBQYdf^p7IDX>!u?6d+qy}-_J#<{id8Oxbqc1@!`OZefMKU*;NJEy?TEwJ+h<9O#6 z*agmFPV8sTSBTf?3xywkX1J)pz8Ck}>^JRVtD-)qFBU#or%MX#QfDzI_OrQ<^K_Z; z)89g}wx<1|W2+PWxW`QrjGyHu3&wNhi>+rJug=~=yIk~}UXN<7D6lDlael89tj>K- z6|CayD#2Pgo0eD`yiay@V%CM%zH0>IJX|Xn*Uxo=vETKARcuhS8w5M39@X3^7@u=- zlVBY0<^sD#u)#jwt%C9X_cp<}&TbcsKJO5W{q7X3;<4N%82jB_VD|{d`~Q0dqtE*U ztMgdy7mVl81A_59dQdQq_fUaN7mVlF!-CQ0BZ6_ic~mggWrkpbJ@>~1<9Lq?#yFo4 zjO+7B!5GU^g7F-iDHwe|Ef`~Yrof(c*4*>81M__sr!>#1K7 zemLgKf^naGMT|lAs*geany}>p!81Z|&#xZXa5w(t6>kFRU| zyLumh+jl|C@qOrgE#~(cpBLB{&eS^l63o^I+q_flZw>eErZ`t$CENLgE%%?VlWj|3 zOTXVF+xZK&-zM7yge~XmyJWkdupMj(qJ5ui7ZSGow*WsR+l7TK=j_L1yNKJW_5KrV zo2vcK$YvnAi8_w@<$455JH5y%5hI$H zpFfkI<%A#JXDL5iGk+yND+oVax63<=^%-101vP&sKPw79tmQwj&GXK+_HXjDlJLW| zRx@o@KOAFa^o!^DDk65~)eL^3?W)3-ew!!T)r2kA(`?Chbz#djGkdaKL)bF!Ig;&~ z!j@|&`&~=ea!<{5Xf139`FzirVrb*G@n=01x9v8i{x7XuujiUSi|!!C8RE8c`8d>b zL)+S!-MJmO2Ir}=hR+u*yf68@u6bU}=f2{4W-VH}E%p4uc8JHYfZI|p=(cgstGMk# zg7JHlg$3ifqG1KPf&(=p;4 z<{q}I=*PTNKeF9KKgKfF_b?U93C`61yI0j`!~VOcYwmZuyJqZrh#0t!_D`&W>e$QL zfW+)Q_B&q0&hZ9{eiiGCHeS@2_m*P?Ng6^Ue<;rX7Ba4KhTCci*s11lA4+;s`YI6`N);7xfV}J@mM$3hR>;p*Qp`ycR!42 z`!e+~k0G9OtksArYxtZr-20M`bj`Ih%6-Lgc|F|UZK+2K+aVs~7`LS!>$Y*vsJQL8 z#5ga82lTK%8}_|Alip7E`pwBVaNOVjegJ&72=ANiJW01qx{~ztpidKX#d+jh#oV|a z)iX?ZPL-aGad~yEG~{ieu%Gf-@mNpwIAiXN<-ls48$RzH?3#YZyXG}=f@^+xsg!G`#F#|Xy#^H{<7 zz0Prh@mx7xFn;bnK`_?vM8TNDNrG{1W=yeU8rPaNQ>@~-pDcWGZEp? zp~u~#9#{XQ?Zsfd!~2t=`y|~r=tn~3% z_RB2RJF{5t%wjH?#d?#)dh^wOh;theb{@2SwSZdXEEy<0Fa9O=*EjcW4jSL#&!cVeVO~~p<`^OV{F$! z$Jnlg&an-#x6p0|i?Q7a9b+>cW4i-7#&$b&jO{k)7@O%B+pW+swp*ZcY@E}(z+!9< zK*!ij$Jp+Njx}I>u%?#&!>MjO}jd92>9UGr(eOk3z@TOvl(BfsU~~3>{;e z4jp4N9bz=ObdHVJ=BL48Y%`%_Y^GyuPeI4no`jCEJpmnKGaX}l96HAK7<7(} z{+2CNdl5Ru_5w70*(lFL$Jk8A*q(!qu{{f&W8*z{uG<^1jj@@IvAqr* zV|xuc#`Y?7jLme6?G@-4+sn{7HeToRdHNn~V{E2lZ0|zH*xrGTvAqo)V>2CNdkZ?o z_9k?WjrZiAg2g%g1UklMI>z=fbd2pI=os6F&@ndCF}4q&V{Gq3TWs;O8TSy|dn1eY z9rkRc$fD28;&m{ycpc0v=9yW%4rXSNF^A7m4#rrIW>=YNoclm~3+;2!kDsl-5UhWF zx0){n<2wARz`hoYKEDx+`F&ep-wDQ8z88$&+5V8&0OV#bYd~#(iWq!MF#`E*O2zQDAcx*jxoRx3lJ!rZjI`O_}3c>btoBQTHlsh^F*T6IV)o+H_3JV!D!pZuIs zn|#LcIts=(I|;^mb{355tBYX#{-CR1{QTZcFy0rfD;RyQ=WH%Z%o}r9zseeZ_S@Yx z>%D<%=Cz?~-WzW0nsdGpbll%Jsq%*N+J>iXj{@5?vAA~BvlN)mW~rZz$6Pi~EUumG zGp?P?%qQ1Q&*U?Xw}oJgb4$TkxG&i1ak9{NGYwXHEX@MYu0w4Yu0uEv|UrF z`$OCHo_a6nT4?&;6WY!J>OG+C`aF*v&|<9ngT+`!xn|#yt{LkH*Nkx#m zbc}TX=i(UfI2T7l$9qpx^mn3Q{H(5?8A8Tq zhQv7{vuB1hHQq^LJjQaeGp@H7=PB^n+TQSX3T*8hW#3bU--_FwR$!+W*ckEpIo&75;62*?f^i)@AlN7$^FhIQ?mQ$I^PMgj@9!TLjQ7Tm2*x-c6|BzV zoFN$JjB(Bsj6R<(uxAA08hKVQu8HRaV=T`X*b9Qu z=Zk_>d>&pBjOYK$1@?-wcn+}Nt6;WgQ@>X5^SWUCO!h{Bz3D8D&wg)#+4$6N3qOqQ z9pQ)UUH3!%p76u#+4}|dL4kcJ7{~jlz&>^s^Jl+Lz+(MB6@IuD?b#;zdh;JE^!vQP zwz_;)pFG#T5dC=mekmA#ckC;{xV~4PGV6G}rhF~>k$oc=eSRw#_rvc5<9_(PVEq1} zX4-PuYlR5k-7;CM@<9`2>=tsuqWW@V+*5zl>kMpPcv1Y%Be$3sTb&=N~ zuOYvReyrzjg7LmgG1g2mu6cV7NREZMtNi|}nTVy8xkCFxFvc>wSQ9*t?U^?@7LNC) z@X4C}B^dju@i>2qF_ypGr`m7+0gL<1zhZpWx8|BzbK-h!<}8ju)*Q^npq@?Ga*WyC zmU<4iRqJ_9;fM1*m;0ffTlispS_nU6^SB@CdEK^!l>*H^&oHk=z}iD|E*FOG5cERO zwa{GW3qp4cdI9K8(7czLA3A>i_-0b$&mS$p(uUPQYD*34hKivD*5_xfcS4w_b2|v6R ztXzm^mE>n-;fLo_$`GTwB=w=Qr7|BW$^z+9um}!j@}h z?PS|t*fQ^Rl5GcJ%Qck!)(Tthsksgvg)Psyb}5EVZmYgu+<-r2W41#NXgk-)HxYh$ zzp<%cyg%5iz&0F@!ocTj`P30@XNi+o~2^vZ(PURQ+5>nxIf-qH)}th zn>&eqtoP1>@!F?+lI`OCV*V8i0c~X3&+4C7(eUh&{#YM5N8Us886(}3HsRsyK#yHSzsrPoyw0j~$FQS`}uP(2rhK1eec0`y0chEuD3I*tl?** zXCy!7gM02-f^qz_1>-tBM=<(4S1|TFFELyDZ09Ey^CgS|O2vzU8kagOr3j&Wud zVoXJ#?Z%wn9G#W*vIab8!5GqV_HW--ppVw{=9I5V@g$5`{) zHOBkH>r)QK#@ZifHwec2_8SG`wc{qicn!ZZ4t-y<07a<5>VBlWx8oTK|hKYo6=UohtPfMA^42L)xbckd2Y2PohWoGYl z5B(tdG{*5h6pV3xBpB=ev0&UoKM{=A>Q4pZ&t*OnjNfs7E*O1&Ay{im675UDcrWr* zVsWkJ-zkpkC9}9*GP4*s-q&e7W32TziN){KE=&F87NieR@pPj{hiS_)o${N1Q`9;{W7QeYIH3puLUrs-G5LoC}LnO7jj$bg@rB8mqpx` zdQoA^?}-<4Tk6H#mOsF`vcl8mu&6)ABLB;3YhIpv|Yh%IcC+y{y&Qp80<(M`X75f9NRvte+z>3*V}s-;=CP5p z7WO8Z%_G~ZEKNw6KGo}jJ-Rw ztw&z-HiouqK}$QJS!}f30n9YV+TJz$_H)g#`nqQ9eV}7*+d;?NwuO$lZ37*1>kS=q z>jfQi+ZsCNwiR^DZA)m&P1Su*uo&AO&@r~%p<`^jLC4s3g^sc90v%)989K(c6LgGi zN9d+v8w?g>8w4F=+ZQ^jNcn6*3;hfmVNdk zz%jEP1!kz?{IHZmWQP~n5zct6v3^`r_RI^rm!Re|E~0J4I-nhfz20oujjf9r;S za}w8bSJxb;i)*gs&aPSePSA0^oQSdFdN~0)u9xGX<9ay`IjO|S57~2`pF}Bm8V{E5E$JkDVjvSqbW?js9ayaUWv*HG zOI@?>m$+u#FNTh}T?8F-yAV3&b^&zE?R@B%+j-D2w{xLmZs$P9+|Gt>s_xUkVr*AI z$JnMq$Jnlfjy3I>t5`I>t5$+G11ZhCQbzo*P#eZ0*@C(N;aHrOt*z zn`5;7*|2Y!s3O3!@O@eW~ z-Rz8?8>|aI=iCC;+I`sJHfXCG`K^V%x2L{$roMNGzPvZPEA_p{`;y-+`ZA__Q{Vf& zFZq3iz7M3n52d~jioU$pou2wW;(f^<7JYdQd^Gia%=?nh5Pc8yd3ikbebW1qKau*{ z{IQNtf!X;?+oz%JeCD__y)Vx+bW2v+AlKP<411mha} zSg>L4^Ao{1-lu}GzMl!k{5}_q{k{;azsK^WV56LUB^cM%*Mjld|BYa*^|uA~onVaf zd%^beIDZg~^YEi!to2WV@!ImUVC?scV4R0v1v|jw{7o?W{Jp^b5R7&CQ!wW6mtdT~ zzXjv`{UaFj`&Tf2*HCkvCEEC&g7e<2z?v7>Y|hl)FgsYhZqFh74D;I0DH!vetH917IUIcTNiN- z7ZrZ^bMVCk8}IY`)Fk^$6eFk8CNyc;2aJ-*B&5y3lXs zsk7$C@sv-Fr=H)_%4@bvVLbIL8pf&E1oyeD@X5VnIl*|1Qa&s0bNNC)^=ugCuAX_s zb7uwNldR$_u2-@Z!QwgCs^DiO!MG+@F0fUc#qrs1RWKW$dNtvPv92!skgefjWe>>rad&?r}XXdt?$l4eBbtte6r`^`J|IzJTE#H&il@yANPbVg3)JJ!MMkE6O8vS>k7vGcP&wC_EY0={jMi` zGKcjAW3APAJnuIU{m51pak4HOihi6w)sHpXNc3awYl`u>A8suAv7X%p<9(K5teIk5 z^Q(zin7hi4HQPkQ!fQwm!5GWhVmzM5OAE&FHWfZuv&{r!KQ$ibPcg=_x%*W6O;516 z-)teqXMML6jO%qPF$USzJ_dC!VaqXkyDjxLZrjYZC2yS5ZG|7Q?c5J_AK{1f=_~w@ z^>aVe+q-QGkHx;{U|zd}wTI?h?gnk&%TVtMT?@^1z6*56pm&Du1kL;QouKWz8tS7i zZ~Q&$j^Oco)*Ybp-pzezj{@6MFy^#Zf%Px20R=YDSzM1CZ*MSLkJS4JTds}nILFME zdJ}QZk?kwu;b+W21va?Ah7{OP!5GVa1vboCtPR(mJ&PrttGztmXv;DBrrHb_V{%^^ zQD7qrY?NT!8}}ECb2_@f#uV6CXE9&;w6z%5%{XDp+H9C=v%Oe1+*c0}@i0gAOcmBt zJ^N+2ot|h13ZIy1&e$2DbR7RJ{dah)h9v6z4}Dx zxL2P59rx;6(_Vc%c-*UxgU)+3V>`_m*SURfJrJ*_XU@Fu3^2PM@S1SC@WVQuDHv;W zR)L-EjN@at|Gtws2Q03^a|?dX6O7lR^9$?(XK{S?yAaIAr@lz|VU8CIKV+Ahm zAAUZ(tibBT80v3DSf85&{h{euH7aWecmn@ zW4XiG+*TfM%6T=_+Zzp&3M<$d4g;1e+RkdJ@!P`6MdXRpzZu%|2qqD-X$30 zyjw7ic~61eo0!GIc3)!lK7HOV7;F82VC?sxV65dsf-%48g3;&0g0bHtf>k`hM+M`W zm?0SZJti1qQO}oQeq{VN5p6sUJ3J1bZIfO*>3r@o)(fj^s_}PXjOU3$tWOFy#K*i? z_~aU89&wLhy{CwN^m(&j?01)7Tu(CuV@=iX!*j1tjN_eI@cE2jJpYwX?msF&?ysK; zpPVE0d+n_CQ(}&o!%V^KKFr$Fg0bH-g7N1EiZQ=uML(`V^*jCS_nhcAq(01==Myu3 z^f_aS88xiyJ_Tl$%>9Mr(->?0qF{{kCBaxT^*isa%rV-_q95~nMKI>4e&?OByej%J zht~vS?kYc?M~YRd8PwFgE_{-`QDAQt*jokmc7eSk7|-{21>^m!`u%${{{DU3v&iiC z@8dc%#qr(~DW`=Rh#F$ZWL3C38zxy&pY&IwZ$`vdJ` z(XUb^H8r0U*r(279_(jx81wiXI^H9ECj4^ke^Fpx3dY)e<*d2KZU^2+e^X@*|F-OF zVas#&TeoFA-w9i?AKaGv?)PpR*9-mrm{^w(XZ}5M`^?UknSCB-%go+q&OfE`j4_s< z6SMDo*>=6k{?@Qo+7sZInIZc8CHai(*Tgy*^rn1rulp_avo0L(cfl&|^AEweNB$`o z_sG8zv+p+8)bD3AmcLU!>%(>VPhz!nh-3EIvA_)RzTn^F(-_CAx!zR6{F!2$%>?7x zZ7vw+VK%{deVAP^-b>9P80TS5!RT`?XFQiHmV~#p@$YHcP~d1f=gV_oZZQVux#Nu3MpK-1;GhdwR`I1j#9IvHdjB|d$IM)jZ#<^ZlFwVq%{QkRwxwPvxV_6wI ze*e7^bY2Jay;gy>7L4mo^&?aLxDHkp>yB&{!I)DU5er#c!A4lhXlo0`@zxQHKHCY# zHQ8P;uE`F9ji?W$rdBY{gX+gMsp90C>?nLPhfdDan%n>^uF3VG|N2PI}cdCqN;d>Z3A-$I;MWIaSb=Dw+5?6+BgZ7vwE@jV6O zebE+z@#hFz3O2%$MBB<4@5}5x#p3pjTD|i%fpr*M*uxGG{aEV*1*6aL1vWu2)^4I;tkpq+u{H+_Ho|Lhh%=ri z_BP`@6fEvxhZSrOPqs%0Th``Cx1~Nx*fNHrliyXln;sMR63}CV zUL4xaVHL~DU@;ccF_x8}V=S$pV=OB|$5>3qSSrvlmKC6zilq%$jKy?}r8RVnWi99! z%bL(J7Sl17HK1cGt3x*xODC`xi|H6kN9Y(!Ep&{f19XhVbd04vbd04PbW^cx2o_^8 z9b?%5I>xd-bc|&^=opLX7|XiQF_v!7O~uj^EXHCw#nr`%Qnz4mfp}YmR`^?7Sl17t)XKqTR}G!%g$gi7Sl17ouFeZJ3_}; zc7Tqtn2xb*4;^Fa2i;UG!@*)KreiF_pkpliLC08zLdRH4$5@6y$5;kKTP!W?fEK@J zx$OGJpJ^|J?f@J1zXZA#nxDBZhK`?UFM^JrX)lDfy`Aw~03E+)IUhQHraceZVryv! zG>eVbg=@h~bL?wek1#{Dt6g)fX|5UjRnYBWM?Dof=5{4?%xwyE%$9nf z{egBkJ~P|B6We81+20!5@*d}Q_=)#f%lcZ2erbD0vRy^ka_;YRTk5vLmVWO_emfUz z?@qQG30vM*opW)0Zq!={Ti#dQ<9@082-^d_p!d2h^{#Hq`yVr7Zf6zl;qQYVo0r&+ z^}1g$`h1|k9u$oEKO`9Qoh}&jdswg$7CPD^&eVE&6fCZ-83o(NlI`QdmbH1pZKKI^vB&k0-B_IbCZenHr>)-Sp(^-FH6?%`hsvptBmuR_}% z#BpBnzMPwQ5C1w?XV|iauemMPV*DQBjpWDnHRk;mSiFaSGug&_cs|F8bMtnxh1>tm zn|HwC_Ym(EY~M?^?+aVz`9ZS%P}nlQkCN@j1=~-O?We+)HTx{telBb|hhHSyFNH1V z|EpyCwXhvwMMnE3*;>8xHKJl~p&ga3Cudz{?>64I==WRq!*yhq``I68T~@U}b-(XK zzv*U<_I-iO-zXUtL^Y~kgLH17?<6q&2=TprM=CJXc^5;U$oW)v^H3y5enoam& zt!5W~cuzTJmDSWR_Bjf+dCa*AW6mx7aD8yhabC+71=~F4JYvj|?q^=%hc#$fWeuMj z<}29F4;E{?fZN7($o*wO;fHg(P?a_8^$Qnl7Xgc7F6y>%OrDpE2|pve28$;@O9)%m zWXWW^l(1zkODEf9ge}+5vdMNiVas*Ae6n3Z*fO_DvRzTwa$RS?t%NPtU#`PS!j}8T z$|;6bge_~fYO=NG?d5ZwwyU|HIN!{Bb-{QptsxlmT2ru!uamU|8|AFEVEvu75$qsm zZ3W|6TU#*Zw~k;H_u0j(7|(|-1Y;~)3dT9!N-(a$tp#Jv zdI`o@dJD#W+Z5Qgf-%nR1RGgj&YC`gG55ZL4flTi1gmqly3`oF%=SL&ecZO0b31UI?<*MB>!1P~Twp^AY^Y$2WxoO&<}BvRehXjM_`8qc z!Vl+gM1hSguu+1I^m^`JV519cjI)>%ecHN+b2wJ`;a)sWFs@7WTsYQZ^(mG>cfN|Ui%IgjMv&D3hc-NJF38r7L2hRQ((tBi}|u&TR+#Rj}v~l_Kq*G6AJ7^ z!C2pu3hd+pJH=Vdi9T&z#5p`w_~AYMX@c?d$zGS+0&RRg@&5XB(T{aHL$Hy))|F58 zJ5%%{J4-O0zv_8p+)K_b^lLM1*8Di0^2zbKV$I~bwDSDUDU8=$^kbZgaep{h_~ia{ zo?yJ6P(CXj=lO+x>UmwveO-|Q&z%c|PqGW0#q~;d5tyx4>Wd40E)k4t^3npk%vl_t z{p!GMeCkQU4`Z7w{E%JleyFbye)zk5Qwr?L0-Gur$GfV)ra6oGv)|QVvHsTxKipe( z7xPbcZK2Pqs?x`rcI-i!s>m6|mN}@V&h%{BZo&3hZ?;2K&7M7U$zlVaxUK7Fdku96ujoooV}a zvOT|G`%bdGv|u|C*Mb-BE(dxfy&wc@>GJ5|{7+VOs}y;|7vT-+Pik{Ivp zZX3Uw;JNq#{N{bQ;<!9u&Z}G%zN;uvU&RW=;XtpLfp93a0{u$p= zuMkSWb-3>)-N% z@tVW&<2?mqQ_rU2+^-;fa=eN&?(uCx4l9DioJ_Bsc+Qcrs(6@xD-jRZ;!4icI&BSR z`z6=qTF`Nwt_dC2=^D^+ovscY*Xe4|ah;x9q#E!1HCU$VmHL<;Gu8DroF}J?ZF}FU@F}LlY zV{Y3*$K1Aoj=A-Qj=A-Mj=60O-BjJb#lMS>v0VxsV>=5v#&#lfjO}pf7~44L7~4?j z7~7uEF}CfYEw;E9R_rabojq>Fb}6u3oyGl~b7jwRTL-?W`Rul6Td~2>>{)qc%eI^Q zmHP(rh-;*aYmU>|HP=Wd*IXkVq2oO4jj`f941|vJFaSEvLx1Qv4|_q!dDs&=&chzi zaUOPuZfcDT2aB-{gO0K72OVP@3LRq`0v%%;3>{+|1RZ1B7dpna4|G%K=0vbq_X)09 z_wlY-_XAzC?gv1}+{Qu2+{Qx3+{Qr1+(tvk-1di#xs8I3xs8O5xs8Bss_sXD#n_I7 zjvSwbd2p#=os4}&@r}yp<`?ZL0fF<+^}apu4CBS(FI$3?qIZ4&l%+B z#$()1Jm)!=#|p-KqvHhQy~y!`jrO=t5bOYFCki&+*-3(pb9S;|^m&S4jOA3pxK2+? z%;t$rJ*$s(IX(5WKCI^%iP>6W%RX&gv5~FIH*n;5XI6bS+&`V?nq!{jnsaluYo4Fy zx@O+zK(7P;w7&p)V`zI>J0H3`j`p%mwHPd}@w3jv0{cyCGHPhSUXfJD~cfiqJ)=ck+qrI$|-U&y0 zSu?#ej`p%~WzF2Lvgg1HPidyXfJD~?XwO)@0cDAZ7*x4ZS8XXnzpsb^=Nt& zw7smE-XBMMSu<_tH_vg?w${1cP1{=MdN*zRh`p?tw!MSs(-lI9C^?@r-faFHWqTL2sK}ZGUOV zp*;bPnHi$bOOnsXE={bHL2t?@=kv1E&$@8DI>9)vlM;)y&d;5UwazToIx~xfwVs^D zi*a6_SgiGqg<5A8eO{4#Mm8m}SZn2zwZ1a-i$13c#&h^8!MOh_#&dC+=tp+7U_6hO zxZWabSf8fImK3a&8K7-&f&JOIA9G(p^y8j&jfjP7@LIu`!*znOpBj&ANiptE*9)KA zi`Nu6klj$|cVmI6+_^X3B>FLjn+4;ZdW&G3zgq=kJ#Q@W3dZ}W8G><+9utiB8;=Xd zep`w8bt%>-*5-hWz*&pJIL{E$7H{5&W8 zur|ARZDKrR&nG`G2tTaz9noK-`mOd9bwBD-c7de30tm<_mk}h z!j^OYVY2;5*fO_|lkF$MmTMvV{Z!cU9M5(5OxW_=`8>t&h1;snDdT)?M8?|hBN)f{ zQut)QL82e){*~y*`5PkoaSm+F(eN66#p8@po_o@2^y(@UXpm%|$UweCJXvR*>?-@+*1kQUo(>p@v^~rho zxe(_*=ojNO9pl^^I>tE=I>tEwI>u=_#@Qb_#<>@Cj+1kHvd^um=P2}xahi^Cj)abJ zj)0DF4u_6$nvQV}gN||T2c6^O`uwF(&x6n}#%VgnIT1R>IRQGxIUYL3X*$MvAasoL z0O%Yi_xDnq$D&`1({zmU80Z-1(a(J#hnI>vb# zbd2*<=osfI&>Y7iFdgGO89K&!5_FD}`@B8h)tFkJ7olH_({zmULg*Oh1<*0h^Pyv$ zremDvLB}}Hh0bwu&-%5nKBuBzjMH?C^GfI#=M?A|=M~T~PSY{Y%b{bOlcAfc*$rSZ z7Sl17>!D*T*Fncvu7!@Vn2xbr107?z8oH@i?f{Fin2xdB4jp5;4LZhhD|C#-bd2Q| z=orh*&`rhi09cI0bd2SG=org=&@q;Kp<^tjV=VVT$5`%$ZYq`;U@;ccF_uT6V=RwA z$5yu{;GGV=*0Lc@jFt@&t5KvAhZv zV=*0Lc?CMg@-lRc zfR3@e4&78NAA`kMOvhM0f{w9#2pwbj06NBEI>z!obd2Ra=%!-%3M|HAI>z!Pbd2Q- z=oriA&@mR%F_zDuV=SLSTP*6c%w^X%&3^O!I`-pU{<~nj9{o{Ze+tI?-M<9mXN126 zVNJ_-3;35q->if+u4LIYcspsQqLi5Ip&!8!+Di?}VlkboPqyvA_?QbL2qi*l!=`dAzS3Ot(zBl60G(2f?0ez~0^$ zI?ox$9xNEo@gW5^R4~?eVNoZJXSwEBxF*$jto`c3C(m&;9+~pVb-tg7g*nVo7;nx3 zQ|p!c;yi_Z!$d4RcR1eMRyuDLKdQC$|9UMkPL(_BIb7s7(sSTiQu&Rj_G^6Z{Qt_2 zamMQ}*V*!7ZaMEGMeh7uJW4RuZ2tlqEf~+0F@o{@8!H&k_i=(T7L_08{Q%LA=d|+4 z^H0@<*90{l_x1zDc+7pgGqslTxiSH^aV<@Rj{C0ZxbK?gd1p)4bZOr`2t4m~tfyK_ zJRc4&^gBebRyGXUp@Q|VM>U5D#x;1jU={CogkYTaBL(9+J4!IFiKCtIdo3G^=gKi) z@$=7d(DuC-`LTt*$EUt0roJbLzMPMfQr}a&FZs!$FJn42^*!DDlAl)Sdq(PeR_c4E z=*x9-cItbs_a#3^^yT?=Ug~>+_a#4H^yU5Ug{kkw-k1EM)Hi-_c?np1%qMLxgO1-@ zUg~|D+oC}0gcs{L39K_T=cLYUxgW>xEhi^GwdlvXTmjZG_`N*Ywgcn2aV1!nU^^w* z0{!=0%T%zqcCT_<#%H#SeVSkuGeEmqFrIVQ2*&;NTERHpb%NDZN5eDy1Y^G&1mk?% zC>Z;)3Pq4b`X!w1F%j_?W>&g3>`$a$I_kdvR_n=_h z_a73BKBo)D+#eQ z7L5I#DX?b??70GaUNFY;LV>;LEauC8V{WL=m-;2)hu@35TwtqCu_`uRzwD=;OU^pI zB7Bm)T41j^i#f5Mt&2E^uM0oCH!8*UhUmvJ-W052B|v-2S&MpE(+)@0_iZq{=5qXn zr#AjQ^o5}9`pENiL1?=^QZE2)*J1t~e12%V4pX;;wreE!{`sKoTFkiTh0gQD-yeNP zNjQ9Ipk-W z1w_Bm?sGxGxDP0w?6;8UN4BtF6(4Vr0$a4e7AvsD3v7u3TT(DyZ45WkE~K)D>{pLu%FFu%wr{J%Y&_z@XH)m7L4cKDh0M`fvx5&#!8>7gT+|a zgpTL;8p1F4v$X`{KHR#%+7wvZ0$W=!#sVl&1mjxmEEwmiOM!JQux`#`{`6^UD(1f)be#Wng9F6&}xLJ{vsRatwPOcdX5}VocuGZ&zS_3aqbSqpK?l|3+Of&gu3AwnKsK=q%<- zpSBj`y4gwCvNoTj+NfuY^E1}YBA(Hnqk85yYpR}2&d*A_2%luTCT8D5u&HNobG+SB zKkLK2ZTH0NyArnS)4n%h%gpM_Mm8tkz_H?g?>?#i=aSTWVto6#WYjH_HWvBxp!Z1e z1O4~B^90 zN3O5IuDOnfLR${BA5!h#kYD={(|(D?+=nG*-I&Ah0vjP1a~mlbeU1{0{q|2R=07?y zi<7k*5X-pitjLb^GzI=zoPN)&6w^R_;S$wf?gJy@!8wUK#vaGG~=+hmxkt? zm|hCn)?JK^KNGwr#kMMJ`9nT3<{0?QT z*R+|v>8;|&ae{Gv4-kwo-ExgNXjr4B7>|0^_;52oTPdBp_8$AG=YChr0PVm+EdNY1 zi^j1Gu(#00i+;Q&v=*_DO%VN9mx+S0-$4a-uwdK+4iSv&{!qcV?hX@->*#Q2@!H8a zj{viClKRMk?NQ10Xkp9R9OJgs#|m4Hd0g^)e8KjFWP75pWi3x~Tk4aAEo*y<+ftt@ zY+38m+?M)ux8>&&t0(h51FStXZO?-40L^jE^uC;%_-_Wz0qY!mo$as3-R9!T##&G`k$W*!QyAGiwd?EC)-PeEo*jZvb{{$ zGMBn!JE>qhIoVz=Y+17_lI;{>%Q?I<*-jO+lw^xQ9J`NqwD^?XAi7IbqA% z+~&5_F9}=L>h|RK^@8mk$@VQ_%Ua&)w$$$kTh{h2x21kx*s|7lyDjxcZmagNPr%xv zq_q7EI__bgrg7pP_9a+n^koe{cU!fGeVzQoJ?vYsxQBg{Y~vor=l*kUzH{3awk*)h zmi2rXb+Fuc?=cWk{#{GGs;PO^?{T9o*VY^G6Zf8z)4DVtw0$$#o+fNL_iwo^^_jw!e&0@h z&neixlWflywmdhkzO+6!>WhRe_nLRzFZHFu_CPP_du~fT$!)nO#F`B+T-V-*pL{)J zy*?0(_4=^DJ`#-ie=HdD{X{V4_o-k!hd*aczB3u>CUGekE*Ko3GuL`Ws=( zT78@Repj&lKH2^tY+1`6-In?%VawY7?6%au2wT?rSGT49&280n?RT*DSV6S?6WaD5 zj`N52<=n*U+TUQEVapo+<+fbI@!tHOef+u4)t9Kx1)&Y5iI61I$C?qu7dU^`E;ombehX7eT6mco{EIDfKTK-hBr z7fiMb30t1G3nyEvcbi+1g-REV1?kIA&&u zJ{M2@B3mLcTmNjz=V*sZrhfJ|$6HD;&g;^N*;;4g|NFtZbRfVnGdnNYGPC!&K9@=3 zMYgP9ocraRabK{0oWB*StYJ@CUf8l06}M$PD+*h(mE4xKYUQ@;+2G#jcV)Nbed8Lg z8UHHD*Q&0^xS!Qs^S*I4=r&*#=dGdby2M!4g5DU~Ue?xx?vA6qteM^fM|)W_-2+E^ zSu?#Uj`p%FINHma>7F>+%bMveaI}{-(_7+bFKec^!qHyVOmB^&y{wt; zg`>T!neL6Fy{wtu21k2YGrcX2_OfPrI~?t0&2%3e?Pbk$UmWda&2&E;?PbmM_Bh(h zn&};Iw3jv0JK|_BYo>R?(O%X}?~J3pteLj|6UAQEO!vppUe-(xz|mgTOb^7-Ue-+S zjibG+nYOjX^<~=D7T1^QLD2THX4=jNo)e~R9dj+49tv$QYo_>(TTmXnR>Ry+4ljvS!-$D(+XN$3WZ5nrU16T>qxWLEFok=>u@I zmo?LN4Y8Ls)8lcpmo?K9aI}{-)AoO&*vp#frmh`!uC)PeZ6;`K+_qW0tp3S)YAYCX zSUa(90*gM|C!fZcLx;p-4(k8zU=FpZUmUMvVljv8 zGv<()`D|r>pmj<kgP?+4aX)r?H}tXMhFHWN9ppL*sz~)VoTHwiALmHr!0{AgUF`pU$m4NtRX^@asxDmDsvp<2 z{XY+RJoei{mwMiV`@CsOflx( z*L^mxpR($oWc|Q+FK=LbF+S_GgYZMPWAd|;@Wa~7f0K>W_-7qtJ10N82tTa}L<*hjXX=kYzu63P0Ru=MnWGvwm^D_7Z-$Ci**z^&uOO{8-)d zoN$c{OtyOqTgJCfvfWqMGKN9PcCfJJx)_pdhYDNH{eH=In6PDT!;|d@Vav6U{f-p2 zJjZh#MhRP4%qk98j_`f>i26a6>`wr29R zjpMB#`f(WlCf-?ywL`mtuJADMbyFJmdO z14Rz>IbN`eO%2)v!C05Jg&frLdO2R#LceZ;ab33;jJYeHWXdPkd`IDvac(3S*I;MC zc#knr)Qs#P!FWEXIJvJV#(7YT=h6BiPS#V^j7<5gR98ez&A}oE_S;yD$2hkTjB#!v z7(X|Ba=i`Q`1>c$(Wasw*&!lM&fn&uALr;$(T{Vaa^QH1u`XMQ@i@1tANMR(7p`m7 zkL!9HF&_IJCh}tr-9L^Y!vzr&;HZK0iNpRz{!S_iCCA?Kvx}v$xO= z7x_)EM>R(X#_^6Uupy!@TrWq7eyp|XM|QO6$MsTr?&}gaRI@y6cF)NU1N`sDSA918 z9OF3G{5ixit{MBWA_jg>d4gaSALm48bGo)eEC2i6s8x&b{yxz0drH&sdrH%lWY77s zxAz5)=YO3^Z`L<5Q;g#z5kGTV+-n%`kyw+HMZbzUMjIi<<60gn7}u71o+{^>YhpYod_M*uIadTP1CQrHyzjRLC|r1)|vEXw#Me-8#r?9{@*x<+62FY zQe63*GE%*JE5??3!!gBIx#D z)E7c`2>Jr(T4;`UK6J;R&x5wh%^19_oeSL=N9uE+yFfGkv!S~}Gv2eH<9wOo`7~Lu ziW#6?E*S4^t`Lm*O%aT-T`3r2nJO4_xXM}c`p~L>l1&4P-wj+X{4j3yJVoALTqDLH zyEct+U15yt1>+udgBXMC#x%xF!Vl|svtYcRyG4vac552rHsOc0x?T9;8oi^+8h-XZ z)zAAlKeWwb-dPy)E-@z82gl@n(^+$K?Tu^(#c51aiud`u!2T()e+#VUR&(C?=XV^>=F?(h zYv#7GZd_~43v9Lmn_Z00T6~LtdlSdE=e6d!;dN^cF$Vq2SzvP&*xUuyqQK^HrsjKI zusGlI2|v7tYbh9Wn!msnD6jK z`>AJ=vQCQ&pJYoE*pkj-PV8svBF^Db!p~R>7p)ZA(xM;7SVpi))niS~vd&t#Zigf5 zyBt`1=&{b0nb!Do`O?s}V7xb63Oar+UlKZgE?)vVelA}eI({x+3_5-;Ulcl@que8w z_qgN!!MIl_uu6fgSYWN3#X017D}lvzxw7yx#A~>UV7%V1T41Xc*y;tghG2|k%>rA? zSjNEQ+YUO$ zwk>pwZ5!wqTW{zXTQBIQV%q~O#3(H8I%rSaXRs2s-ArFLcapALv@xaZT(E9djE99djE19dqjs z9dp|YI_9<~bW?R73l?J=107=<4P6^z+aEf{HVQh%HWE6HGgh!OWIzy!%!AMpyjaX+TyxycuDM1! zxgP7zJ3_~KI2>cec{mKZ7Ped?heF4BI0QP*!@I~qF1b`*4s?MUbt+Y!(fo2vV{V6pD!xMtnYcFnq<<$A2S zL^~5Y=5_{j%V6qmjO|kB7~3V# zwIQ~Pp<`?pLC4rGgpRRY03BmHA3DZ%9&}UZ#u7I-J~!$Lwu>j*R?y}*p7Xq)oaBBY z^=Y zx_a~Hto7b{wbt(4|NX0KSJjqt?(O8B2JXBD?)(Psg2Lsqg7Gd4F5h2WWOTWooUh~v zp5NpMcd?0w{4Z(XE^Xj0Yv3+79I;%{z+G9mTpK>i(|2O4HsD1|zFC(s23EDX%8beV z*VPT&H4WUgh8wD77IK~8nB(gkxEmU{8w;27#i!O{UN<)xU1~FTRhz|29rOC<9{y$% z4>`K;wx*`;JGkd)FHOrW#wU5)YB=m}GaP5`?G4-=4cwg#++7Xa-G$3@K`i$Kr?~*% ze1Q`0eD;v^3gz3H*rh!{G2JVRl?#aUCeDSF@ljrv-<7aqT>n?8g;nS6!u8RoY$;_VAH5jQmU0;!! z8v_yj)@7fMmtNbe{QgYD(mlN{-TZHRgt-@zns95Se{aHDm9dCl_e;nfSLnCS`+?d% zTh&CGaL*Mk&k41Bp^fXfHau_qp!;Isy2*JF&aIb`X=(1;CFS^uA>uW?A^UaFW%Z&M!;aZE`+l9-q!o3rm#z5Y!{Pc#1zG=SRtL$V( zZ0}cGUIQOgoa~tA4=XO)eN?!d6E*xeIOPQXr1GONi0#t`?z6&mmqCce$LG8kcD_IU zB6Rcp@#lg2gg*D~vxR=XKb|#kzCWHNaK1mDIdHx|o+)s?Kb|pgy?2J{r4YWKIoF1B z_sfVY&-Yh>Yd_THYs2yP;%^K$Qm#V2HQeZSs^vSwVaN9wv-V}e+W+3zk?#+NW1if2 z<8|z%1_G;?PKQo*g-FaI6rW;pDxA^bLafP)qc7fKb)gSnfk!VF3(pt#Y?cH}TGsd*WCwBVY>)I!E zpU&7Ze`gsx=0Izv_Q`nqKJ1!fyr%s0-PpC=2)z_Cy%~>l^L%4BwB4*_24ly4(uKwj zpDsV{r}X{Vb)3}9*)c~mnpmjIOopS@GaC*&zLPqyr4c3fS&SWZnbmOQ?#3fO$FX*2 zGd{_^)o|pX?}4uCN$#^7JNB#N$j@=CPko(Y%IZUC(DY*6#dT)v(V=YqlMUJTr$LAPh6yO7bPHVYSBa4)0Fn2S_? zdpGDVTIu#Ny4137(FON2y41FR(FG4My3~4L(FHG7ba@Y_dXo3z!SxP|?vjCZkIpzt z6kFz|Uz&@hgX^DsEmd?m6Lc?54VS6>^a+l;^s zbA{mY^>4)n-IXfcm5nYnTcy%n)#wu6YL)Kl4Z3Sox@#I;YPMFTySC9~4%ex4*EPD# z|9X|~`bL+1yFsO^de`R^x*HmfGkPP#aYpO=vTL85(dWhAk5a%X_BhF0?M}B+c3s#vLCqZD?d9JKddR|2d?(B zqw&MsIX`f z-93#iF$}46_cFSyi=mb7Fr&-d53h7b7+rE3S?P{4x~zrT?`WetLS+o8>#(=cW$)}0 zT%OZ?i*7zUh0#c{svw%b8E%Av4LQni_&nNh?6G4SxPKdtx*S`$e1@|=k8k5T?(vT^y43K5qDwp{ z8eO=Pi!L=jspz`zWk03paz_5QVB$Zu@^xCl!`jQ)@}Gj)1E&Y>8C+}O&kC&dMJ#6q zUN*3r1t4V}cj9yJ5yi)XP65*AjSCa^@61|#4c-83DB*Lpj zuO<=Jce1NVgx82(O(MKz^lB2}wW3#(2(KNznnZY==+z{`>qf685neBPHHq;0(W^;> zH;7(MBD`VrY7*g%qF0j$ZyddvM0k_v)g;23Mz1Cj-Yj}GiSVH4)g;1$qgRs%ZyUXu zL|A{H;cqm;TC=QMVXZCJm$1$g&Jy_!T==MiU;u+~27Us!9O^)Eaku$n|z=L=_yu&y&) zYlL-fadrvs9av2wtnX&v`x%7m`@B}Y6r#P>Gx{EFsHNrXqPuv(>X?_gu4{jhd6~ap z>NI!3)M=iAhZSG*6-?gq2JRIc`U?im`+R}Gd7lgCeJ-5$xv=)QnnYN8TumaJ_qlN1 z=fZiP3+H_ZyP~W4!L^_45KxSm8-3+*;v#6F$F+UliA(3k*k{ z-S_{q7p}irUOIn1(yEt2E)1W&qHh%lxyX#c_!k?F@!faUGrs#S|6%RJwe-DQzjUk* zA@)$sEhe(aVrI>;7u@&Dvkx6dewUd1xb7W$u?FrOC!G6!Yht<7_#_thJ>1mn^D<{$ zNH}_YHjU@L-lRI{o8#{8i!f>r6pDPW=T3W}{<8$Gl%)pH=IAO)G!n$^Ltn0nZc-UQU@*{`sj2-jm<`%p8uTj9A&t2;2z7LqXIF2}P zFmZC$optfw$AfdwzGt&%8B5)c;XmK|?9fmPzU#7UvS*Lq5;^G}27GJD zah#;f+~3~Db-aJRt=PifQ7~)s&f+VNJEHi#tLTF7HoC-gPtgV6TXggFy|wtg&v5$} z?ta5z_kiJ;JL;FmV=cMwCS~n5edpyMGhf&_pXBhMiIW^2DqOdcuR3b>aBw}+>yHM` z^ZrPK?PHbg6P4}b#+Ld#S=l~aY~i0Ww#58QW&2#Qg@3le_W8>8#me>tW6S%?mnz#= ziY@%h#+E(sYGwO+v4wxl*y_EdcY}t$b zf}@V_2dDRJ_3&%0E)pE~51$05_p`+Caits7|K9(9 z8eCqxpA}u=lP9`Jz%?&sn&Uz=J>J%4HAI^O5~ z8d%&Mh5tS9!U_LYeB~OmPW~_)_nLn;aDN#NpNglRH|{l^PwLV1U5D;F61ny|pCd{w zT5kM%EX>_U|d}d_M{2))0I9HZ!N#buqE9o*6G+&ss|kns5F7e>1=Mb+M9v zR}&BC8}sYp=@xeXukqkFuSM#k@9nJD4RsyI#LD?PuHndWyasN3!;Mr3Arlym_1)cY z?3oD-M=UNs=4m2h$DVOMS>LWMoP%yW&a;Wlcq7|WZ<)k!ZJHJiWx z%o3cgiTIhRU~-(XVCpeL!T6rJ;9`dwrI1+-H>jOznayyF(P}vM%It>2ZVtm? zH>cs)V{;kq-V*2Bh9mcR499rukCB7U=V;qY@#?(B4xjTGj+)JHIR36W-+6x@k2SP_ zv11=CXgKQG!*ImX({SXlkl~1BVZ-6Gm*I$I5yMgI-iBj8ELynkDo=>^BC#wHT<^f> zE*?1Voy7v@y)!Uy-a7*V_X~UWO8>xl@AM0t_fFrydGGWItTn>lKk_-TYH-5r+f@pt z#w!=h*eexGE-MDkxvda5=eB&{oZE7Nb8gE9?w4{~CUDMe>A*R+r2^;NmJF=iT;10T zF2}ZR;2hgJfpctY2hOpr6}VrDZOy$Ob|{!Nv3VK1LxZAA7ka(?iV=McHh9cw)+I`7lGn`@4&gXqXXyK zjtZP>J2LRt=Kb*Ca%}$yoMSsIaE|TJz&W-<0{2U?9UM5vc2M9P+kt^|YzG7$+kQJS zxV+y^D44n*Uodq!u3+kZY~Y;RzXRvojtQJ|J34UA?Wn;0Qf~hWoO3%eaL(EZ_TRucw*Lgqv7H_`$97uaekrz71LxRI37lg)IdG2cq`->J?G1g0xbC5m zGmCECn|l}XtOo9E!;L6*V+=R6aOW7VwQ%P)aOX8}=QnT{G;kL-a2FYl_sUCLpnyA{ zXRO5c~O&IazT2JUXd5z9Rd+`Wa%`C`|* ztZ(pr#t-}Q{swM|OSI5Bp99#r??a|e4;Y_t4>oWQ6)xw5oz_L3!-tKZ;bmVm#rBA? zV~j@)$NT)H%|3v8%-HeX^Krvr_e2Btq~Vy)rwqsES5F&mRH@4|h9mc93zut8oX-WP z+Jm2O(0!rOebMMro0p0%_+_KZn6FfRUv1ERtwb!H-Y>SyO}{i39|hMx`TDTva=)YdCTjR`<)=?@ zHcJNnZute-Cv9@^Z#q5`1H@;F;3ywCmNkN#|!7V!)6XHpJ6jqy7>&-qdLQ8sm4_P z#5cTQ-^^<4$YVCc;j^`Yo855C?Hq<r@Ac;ioZsuu9awR@&$bo| zE`P={uwe2UP%vZlFL=0IhV%=ZbL$&8=hi21&TY}aIk(<{b8d?S?jL>`uUFul+roi! zZVLrgZm#aj2bW`8E^v-**}yrrWdi5emJXa_TPkq>6x))4b8Jfl&ao{XSh2BRgR{`mM7r>wJaqlh54^;`79@>8@DmZftbfTPqb^@McDr7*?+QZrPx_ zN~ODv(PeL3Gp0Q^@U}*my}4@f3*N!#j#6ntRx7&Tor^B#p>)XY$m@kV?`yc#!%xjo zuQd#Z-I@*DT81P4wGBtU>lkj7!V6i~aO~mr3YXU*ajqX+r?u6fyJ4lfk$m)dS!bivyg-BInOZW&Z`!GnvgyQXg&oX#M0 zw-2l{h;g>1-mlLoX8BIVa}pu!b@I(ZC(vz#Y-R{nK#7a%2PdufpYg zvD332lrQ)w*Nt9H zBD`MoY7*h~qgRs%ZxFqjM0msK)g;0jMXx3i-Z*+SiSQ=Tt4V}6jb2S6od1u#@Su2I zO(HxvdNqmgw$ZCegte|&+rqkra2*lWdnn#V3G3e>s!4=(hI5V!>$;*Q5zha|UO4|B zd*R(ePfa4cd-Q4&VXZ^fqOkTh`&)QOU^R*GUeT*bgoj42CJ`PMy_!T=`$J74JR*8E ziSWqi)g;2BqF0j$kB(kVBD{C>Y7*i6f9!?p^S4zmg`CvDo!r2k(!ibCz@65>oo+bh z;6H|A{oHlEGUCmoSu8rgBPPeGjw}J z4}X^N+gkL`HXM0ea)FdO=a1bOV+VIm19xuWavs=eesdn@2UZ@)dB!hm{Q|?W?k}vk zMQ|AXMGf4=6(=*!l1mK7IdZAtn2XCAxXTU4cvl#XbLmRMvB$469Q*s~!gWpc;B2|3 zjq5l&t~Gu*N3Sy+`Ci|^-B7sk()i?iV;k3TR^L`Ies3z6v-9SHISX$snDuc>8$Rej zH`eV%mwfIin6d6Cn6d6Gn6d6In6d6^!+zh~TXA_U-dAxd3o+bZas5)w>SwcLpAS@a zeH-jVv1cBv>@xR|;aE2h8;{EuufXsN+5ayneF#AovXuem^j^ z(#!7!CMGcRCH!vi%$M*xfpym9*hcB4kaMfpz6{;<+Nrjm!a24t0_WI151eEBEO3rZ zILG#B;2hg0f$MXanm%gk&H1h8z0`J$duQZSbIiqK#wYvoal^IBRmc;DV?6h~6@0p9 zzrgW~kbIqm(=$Tq9H_IN(^7Nf=X}CBpF7G`h@RO}`y`feZ}>aM8O*anat_0#9`dBA z8QfEbgj3*=X|!dW!KX3w8;Ux@yvL{IjP}@a{|M0&+VRh z!5mF!?BJd;aWa1s8$0IcS!2f>xf~eJanxmUGahs6>_(JYySlKhogM3XYBL^o&zbzl zVSHo9{JFWs&ONh)@t$uO?*+rL&R#SevCM4ZBo_C)54hQl9dSC2HSv;(g}HTh%)`sZ zjy3p-;aF#{8V#Y}RFsd{OCsX>^J2t4jB4qe~3mRJz|9 zUDm~SmG1XOm%0C;(*4oslG{&}?$1V-wNU&0#ptrf>pJ{ublE$GB&TS0U>eJ8I_a;M_Ath^2|^VsgM|SHrcoMcLBQ&2ZFZ z(}o<}Ge#J1y9T@M4ad6P(s1PNe8M@Otodz>PvYFwaIC?>hT|S%98)v6aSb=R)Y`?# zx#BqH!Ex-PolKn6)71>l`D9IuXL7)9H!~h_4mTWe?qN9IH%xe)3eStpNDDkV~xeMc`*RY$gEtZbY5+^H|x=mCtXX?ZS^O?+~1v7V(6wG%@Ox}h&*3P0S z2IAs6*}s7kMg69ze2OEVDJw32-go-d((H((w_uT))64jrs`8n+sVlBu@u$xDWG<(v z?6S{k4ad5g&T!1t^oApc84Smq&S*H|oXK$P;h7CLy5v8L;n-8N8V;Yc8ICimwSk-6 zaMWdv!sYvB*3w*UT*u$}<}|w0aPFc@Jo6Y`xcQ1MHJ-QVy60rfUvxQ-7Alzd7pQzK zSTK9AXTj{j9)WwtJhYa%SR}CiRzxhl0xug_O|o#{<)T-U2rnPKnnZYo=+z{`D@Ly- z5nd^JHHq-b(W^;>SBYLtBD`w!Y7*i6w|(K&<8?KO@EXypNrd%xeKm>jTG6XXgx8K< zO(MKb^lB2}b)#342(K5tnnZa0=+z{`8$_=r5#BI*HHq*>(W^;>H;!ISBD_iTY7*i6 zw|(LKw|(J3p{FJh9vr=zM0nfi)g;3DT$|6fg|%i`x58RmtS@1mC!8t5+7IjrVXb4< zvT**}zVNQWt4V})ufY9+u>P*ECK1+JW<3jQEwUbkhXhuW2=5iWnnYOV5oeOH);{ZB zSZkm4FRZhM^F~lwfID7bgg zbs#^Kp8GOX_-N_|NBhbK$(tg|)}kB*J;03+H_iLmy#nnXD7b78GT)}yfYxSB*b?{ne2&xP|o7tZ@!ctrS7lL+U1E}Zwd z@aWJ}lL+U1E}ZwdaK2v@YSoXBvkb?1yXei*=(x`m;_T{ExUPD!h{lvWdS-jickABQ z20Q;g69lepWFCZtn#~fLwxR;23*@5$M{PapIU#Cqc7FZ&T9tF zeWyBS#MceqRqnp~oi#fAy1$P{{TVOUlsYY0)oCe{)4g&Xvb5p2XE{y>bCqwaUJ6;p z*m348Teut>+;YKbZGe|Iez>Pu!El^YE1EIjRw`qFS2ntgu}aYeuUd4wv}M`;gLSi- z@dLMd@dIAN_#x*tjUTwRiXZUWMVE7AkqC*rHV&?LVAjM&ffr49!@zw4bFJARaNmU2 z58N*>`+dE@{S$uv^3MMsY~A1o22YOb1RmS^*u-$0nVTA}RjxudD_mYX%+2P(E!3{o z_Hzs4huF3>e&DtWu1E5-weiE4?t7!TC*Q`50XL``V{pS5+ZvAh;O)#9aNAd7>|p$G zmh5P_(WM4E6)x8tZs*`~t#&E8d4Eu=U5y{s-fnGN$Njy&pSwQi(XGeay_T65_4>cIp#&z6F3~kUI9$c>Nh@zWo%XMj_@x$DXYU4W2h0zVV zdk2@t+^6W~F?r9mukphk+OP7nztN>82UNNT8eMWZsM0;y=(7F~sdNuDx}(b)KdjRI zhtVas!zO#L8<`{TW6Btr zV=w*NaEq1N9&5PP!X0Ngu2;t!j{HwB966k5IPO_aG8}g9`^|~vWMjwPb9Qj27(43O z^qu?~zCic@z4BVIg*Kn<6oHHtJ;dmYH%!-o< zJ|AkhFFniH;j`O~^3wVKnf2wq51V;&-)B#anVW8MU8Jt#|4rS9XF?Mz=g!$C$KA@< zx0u=xtNR{tavWoP!cAggVeZdquyZ-EZXL&X=N6yb541*ZbCmpkkF!PV1k5uLa{W2? z&kH~KT5~~Q-TT9zU&e8^7gn|xSGE@!TjIZ@vc0U>!e44^iRtpn_R3-le?^1sRh8{E zmF?BWmV1h8E8FXfE&O%Hmbtv4vc0L;!ry3YxyIjI+1^@g;cuyIHGkCcw%~dPM)!`u ziw0)g+lwvhs9&1T+*zeUt7zm9E||u{ZAzZa~t# zuhI?bfA4J`2rjSP2a7K8NtZf2WVm5c2zl6WoCl9IaE}^}oF6kB=howf+qylJmM08H zEK8jyrOs#bm}2*&vE%O;yR<2Fwj-9OjNQG(=hKG6?is@o=d%smbB1HQ=M9I?7Yw&; ziPL=#K5@Qi>^PTSG8}8d)r>iM+1RlUUTNT#FxLp|UNv^C-Pa6<-Rlk98-`UbU!e<)aJvY3;xLH62r%p-%lEJKdp2> zGrH9B^P&s>!st@lFN-etE2B%Tzb?AqZ;CG0;(X438(ik&gcAh#g@6r=lqYs z^$%NW_(Rd<9L?waPnDm1&i@izKIea~bn`jCM(l+FX?^@!bQv?K|E-<%8ura^#*RFG zHymr>j|T2f!!fsi8IHMZxkXWTK4+PaE{0>@bTu4tZr+f4H)BWc;}{N~<2G>P8IIh? zHypW7U^sH`Za8wEuyCx)Jns_)m(Pre8+0eBbSE{s)Mm1x3!dEQQmZK{zf(5oPF3kn zZFH&SG({IYtmQRv z4QDO7?B9XGQNz~CPd+o|2ri!)vsb$L%-FR$Gv+M1jG5QYUJd(ZE@MX?a~qC4=4s&O zH5_w0pW&Fx`3=Y1EnqnI&4PvNuB1cq-yHe`r+MW*XVJi#S8(sZngc!`TO_dN0Ng9E z=9+t6x(uv6&w3is8v5gez>mvENw?EEw5ao-#Y8E&Q}jV z`P>~cX6(9aRJucrE_1(T(FG4Ty2P+n<#%L*?%I{^Xrs$seCV9^+`#)7UH0ZW#V>e2 zqdTglUDhqS-~)=TJDUz~IGffBKbn_pN4?fJ9QE3uf!ok<4hUV~=cF zxV#RDbF<*`+Sdt&J`*Y*YCi)Sx@K(%shRQp@d%E_i#ROKo>3 zy5Jp+F16mN=z@1Hy6&33OK|y`zFT0OL5#C&v1M*_-cXl4g6kic8tz_nS&IXc?w*w& zoonR1S8zJl@H?c^)%nloJv=i8-JwO7F+=r#e$HsvH^Yn_c?>rkj=CKfSZAiY z-#95a#mfE0i3MYOLcxr6e8I$iT;QDBv4L}L{|=mUJ0@_>?dZTcx1$2*-2N3f=XPY^ zfho6t23Bsa?q>#=^oj_tpJb8Pjc#dcwE zIkpP|=h)5m)FF#1+yluDVQ~J zb-}EOs{-fTt_+-WyCQJT?ef4mx61a z+pU3fY_|l?vE3Xv$97ZT9NUe7b8I&R9++agKCoiT=TfU)3b~_!yR(72t8ngI)bqD= zKY`p`blq8S>X^>+eoDd2`^g0}?$2F|s8AaJhj{eg3B z?+cu3dvD-@Egk3FJ%PtI@6QC6V|zMqj_s+yIkqPQ=h&VIoMU@DaE|SkFD+>1(##{FmR6TgTOhq_XFqH-V2;#dpB^7?VZ3mwzmThOtHNcSh2ai zap8@f_r}LXH}B28WghZL1NW)nMzou?d}g?zh5Ouat%dudf%~$7`>KKax`F$qf&12Q zbCft&y-2@wKF?UU-x)jp=H{Mx2FEkc@;t!lnP>HQaPIkM#PYp~1@4E!KpBlKI8@OK@xL*sGV{2_2t))ftna}p$j3536_ znPK&rL+#!3!x&>u;}bhQyRG&~-RClP%%7g)R@*TLS~E4rczUK=%`sk6etN!JZAZQ5 zHsf(_>RE5K9s6<~W5;JWdhT0ohfkN^D5;0&8E~~7HFI{%(Yz*>k#ZF>pW&$W{D#AB zoSS6b`5Hyd7BF_yWkJJ{yBm-E9LL)2VSK{%G#qtVsDWG9aO6If$sKNL!%@#(#wXk& zhU0v2c5u#)T01*9XU9I8%;W$!x#5_j-X;gQMGeOs&3NwL*AnCPF?RUuYdGeipW&!= zf5Q>W0K-vhXUCd%ak9<^8lP~B8IBwlHyk;*nqfDWnIqy{!uVuOENM7u?Z#t{dYbXz zmTK_1wBfK@rh!}5aGW8_8IE&zdBcs;(hFI^aIA?H3zzSAh;yaj^gaT-a)a(FmF}uW zmwK;Obiu0|UB+Ca^1Eh(?pl@Z+D4aJu2Xcu>l$5ZyI#=+uWxjz^#(;3ykXIGe{ zxZY7;bTFV<* z{^qpgP3@nlFgIIOx}o~NwX$_^`JQ{52HioG?qH)!p4(Qs+ZkQr+rHA>p+R@YN_Qut zOU-t!bayej%;Bz;?ruhx`QN?L-NWdzZ}+TpRqy(q1>GTrzVwxmW(iVTuVk8jT55AGm!hKl>R!)MikrO^yd{|H{t+#t-#BuyENA+(DI}gN+~7 zl=A~u`#Hq;VeXtCxZ2O5#t-XwOH&^>+2#2<%=lqV{G)KWK5&OuepGku$E?d5Iik}2 zr_m+8BP-p18C_yHs?t5$=&~-3sdWEsbea2OE8XLaF1a0F>7HP8Sqrt_6OAr=yspDZ zMwh*FauvfVMVISs-mm+abA{UPW;n(;)%e724`WB&PcwGR-w zSBzdwBD_-cY7*gyhgmPCJ|mUdNqmgTG6XX zgx8KM@O$F z5#Bp`HHq*((W^;>_l;gnBD`PpY7$|6CzF~)_<-ouB*Hr9yQ>RXBpS82>DJEwHs!{^ ziv~|_Hw5kzn7O<@aNmTl3*0X-wYfHM|AemzJRmS@Q1y*d6s0&eS=-wZk zFk|0WFt+y=%-Hu7Ozd|D?j82vy8`Fj?hKrByCZPU?e@Sqx7z~e+-?n=bGs#Q&h6&F zW2^fU!R6Q<51eCrEO3tP(ZD&jM*`>A9uAyidnj;@?ZLn~wg&=_Ew&ef%dx!>ILG#U z;2hg?fpcun2F|fP6FA5Abl@D@Q-O1APX->_ns_g`ye8f)m^JZE!K{h53uXB*3y5`zH{J%94IM2h+F;%O=Xv-xaGr;60*`Htbh)i_Y%PIvY=4Dqj_uFDIkrCn=h%J^oMZbfaE|TQ zz&W;G0xLFG_emp`yf-H)n7U6~Fm<1(VCp_$;9U3afpcyX1kSmQA2{bWUf`VDxPfzS z;{?vRbqkzx>l%1$b)PP{9NV;kb8OQD&aq7$IL9_s;2hhOfpctA1kSNd9yrG~S>Unl zjni)Gyf>zA(ABez^FBH+u>9sdzqeirnW6Z}9QTMbR-9r->gT{J52WVw`p81gRQc4) z+#}9xIPMc>sW`1cWY)sv^N0J!);6x=JesZYBOgPHquC9|_;VPJH9cnoH<#fUZ*Ifk za~{KSFFtR@<++}(;*=Wq-19eZ3mA?)y$pVXvx4@8;rHfwesF^4(Mz{>OGtH*D`H>JM44@BDI}fC!S?1 zyUZ=@5+$xRA=M&EPf$=j?~3*#JQe{le6wcbJoGRXUOpx<&x+8y)MM!o^?ll z`#0q0#$#QtZ*qX!z;MK}Vd30;(8g_C$9>R7MwhrZDY}fgsnLboyy$XovRTp1YXQGo z7>@bgvf>7+qoC95#AtrHHq+s z(W^;>H;P_OBD`_*Y7*g1qF0j$ZyLRtM0m64)g;1$qF0j$4~||zK7Hod3Qkyle1k5@Fr*ao;Db>$aLiSZkT}EUbOO z-Vh!VSWP0lSM+KUVVy^uNy1wDtbbvxeb&FQ&Kk}eVVy6WF~Xw)t4V})ZgF-A>z+hS zBD_!ZY7${xN4lyj(pq}$wVv^NkAe>_x{DXAj(M5?=JpqvmxT(ZP74%Fofa&Zyn7Z* z-aP{M3YX~j3Y_=(!hyBN)g;1sp9|-GF04JSCK1*iSCa_meJ-5$xp3a+!g-$y=Y1}$ zJ+3AZ)*e@r2y2h4NrcyqUQHsbJ+3AZ&ihrq&HTumaJ_qlN1=fYa+tao9pb=JFZ-si%3p9|-GE<8Hy z)g;1sp9|-GE}ZWdg#9>;qx%V zF$e!BTvrB*9`5kqdL?{BgP(sk_&Kuj^Dm>zS$tH{1s`2>b3KUf7~==--^CC3SmS3* z$@w_Lacw)^aIEhW8n_b+m+Ok%Nx`YE;FBBtoMJd?b7}*3TH*5e*qt7n#s~k$_#w9c z8b5Gn6hGiIjUU$hSqV@u1`hU47a$(-SEI~$Ji zt}#BT@wE+{^T|EmV&t^ z-)!Q6yVY=`i{INCxZ4efPd6U-Kh7ulA7k?8{l(}8&iQ1G-eKZoo_OwG-V=b-=n~VbMHl>9(apK4Hj>v3$Nl0P71!Ek(f*@;uAg3|_S7$r$GUhZ?qBMA z682Kl^WuJPn_qu6HTitAFwt1&(@eyIP)hT{zU#Eb#=X*I@Y#t*gn-1y;4{GyHP_8g99`m=No5$q+$?wJwd+3kK&!0w@n*3Giw%q>rxhIz{!R7e68eP_Z zw@P;$qsuxTx6&QY=#ty`mF@&am$h2^?QV2ge{~%uG`gH06IC%xY;>vFB$ck7Nm-w@ z=(0BQe3SR2WsJH|;rZgO}Ol3I!uh`Uv8`AEwWts+VTEj6n z(;03|@j1QW_`A#uh9mcxFOY-I*Di8*&;Db)8I8|-ODr=Pj(s+>;fQ4x!%?$Y4aZqG zo8efat%hT7&TcsN;v9y{}iH?85A zqvcHwa4Q&&IdadrVZ0TM9X?kw9P_ZU;i&a0h9j0$4M(k=9ediv$vRuj_=H>CaOALt z;mE<&47&x*91-W5#wTlHEyJ-l-FVD{MiZy1Q4piv+h@G|39F7Y0O640~2%N(XMWhJ7={ z*inhZ~QR5MYV?IV2j(sz#fg5c&&da?G$GNwzsTt?x{pYBJ3H3Ai<5PBjPVKgZ^Mzpv4$fDS2OJP zF>^$m#~GijiQ^5&-gM(J502yCJ)U5E!kuV1VmZlhjCZo(ST9GI+~H0!cKAHiaM+#J zz@2V5YW5$)ajp8V;kZ_vVK~<3nT6v%B44M^3NGK{o!y{2rqVsf=u(?=i!S&)qf4#M zul!!npnGAZdy&zlmKPUY@Fhl<+Fn|8!Iv3bYJGXp1z%Bg-96rw!S#-UqkDDWe2;fk zv1M-hg)ViuHn{$Qso^z6*WKe?SNX~JcsB%>@A0m$bn`tP&(kD^8>=z%+R?K;>w6OH zZZdY{akJsDyQP7<)o{%1ZH8kmZ#NwCafjj9H+LG2I4^9-{Vroi?sppwpZ7Fy_Zp7e z?=u{^-)}f_f533${$Sx)mwDbF3ND`+4>#yOQt3Wwbg9i_MHl?I(WO>TRDPdq(0!`X zecI?!%V&x%_*tV%ZJ#T;;OC7lwSJ-Kf?q7U?#y^8xO`^35;&h3FBeOvG7x_<;G%viq{jO}j)GuE#K6Z*IOq0b;GEkJfpc!(2hO>D7dYqkZQz{SH-VL#tNVC&bdGJ@z&W;Y0_WJe z12~-It9nbKkw_f?qYd_?@uw`+9@!M3wHFMwh*@#d+bh}XS=*dmbCm!3`G->#$IqEf;;qW3T$U+K1SAJ({(4Do?oz3V{%hsX`p55qD+c}CZ zcuu2Bt>-Gb;JJ&gyQa?*T)w8y7g%QyuWgnO~1l1H_E+LKSD05&hLrK`OVx&hkwh|-;4n_z;Hul95S$hTdZ*1+f_T=Z!8~N zz9ughIA8ac4V)uj$ zE|{9HRWM_%S@7^O_8Nh6ZmS2*xvdsB=eBC#oZBjab8agK&bh4=IOn!v;GEkEft8!9 z`=-I=*ft5AW7{}zj%}mBIkpW0=h!v~oMT%*aE@)gz&W;c1CK4X!NKL&1_jQsZ4)@h zwsqhf+g5>dY+DA-v277J$F_Ok9NT7r$F?T+3@)#UJql(`>|QWyVz+`>6T1e^x$P1- z=eBd;oZC);b8b5Z&bjRnIOn!~;GElbfpc!#1|D18M+cW<8x=UmHZpLIZA9Q4+wj0S zwqb#DY(oR**!BvXV;d4!vE_5ARWF6?-N5bB!0lT&cP@T;OXvIA{fe$T3;sB#^Su9F zF!TOf!OZ)w1vBrz1kSZREXK;UJv4Bx?ID44Z4VBdYkN@OT-yT!=h_|+IM;Unz+;>D zV}i@E9UVBwc2wXT+rI+m*p3XGWBX^|9NQ6rb8LqP&awR?@Ywd-e}c>V?ev1F`)LJJ z_frd|?xzIKxt$z1=XO%yoZE?kb8aUD&bb{QIOld;;GEmBfpc#E4m`HHpBG$??cBgQ zwsQjK*v16Tv7H?_$97iW9NU?Jb8Ke>&awSBuwrw2C!F)i-_FK2 zKFMK9!;LJtcQ+jS;8IgFxXTPTy7+W)^7)J7m<8u;+~_<98G8J;I1%nGJi7|JLc$0W5*o192n1W)MXYk9&_vLMwD8+ zy0EUD9qW2_GahzVnf%CMYGcRzxw*yeceCFa@9Ku}t}z_z>{`PS%lsxzVj1Tq70;hb zy&CKmF&sHKJ2+>@TDs2U0C&CNSTAloI5!@3>2GpiAI;uSmm5qh#CfCPsI?o9x-@Y& zHTayy#7TbcdEvw|m$73%+-%0^8c&5%ZyK*?lykl?y3CTYy41~4N7ftJaG3_e(pDZsQ&|n%YNV{XAm)aJ^f{)CW#>dA=SsepnNa6)x8Y?(xcx>R#W2vPPb$ zbe}Z3#P?LC`?S#|hG#0>XN@lF;<-xqd85nRzfkGEXmrW#rAqf@qsv;T{k~#!+2eH` zUNySxo!6=uUN5@&9-nykH0KJn-`a4D@rLn<-5_H}-QP5J%-^=gjycepsjst)x4p4r zyr%qiGpk2V~+JD+gQCu@FRg2gY+8bvfRQ$J{zQ&Ma3K*0r-^U7u{m!|oH4 zA35;s`Fver{@mPR=bkH$-KS-|?q&VxGeXYD?}F(87_f7b# z!2JUAnc|m$`zQQG-~oY&^Yg&@bG6R`>pXRzBe&ey`SZ-bLRYafr+*gA*nbpEF24uP zx&0P6=k{yhoZByfb8bHe&bj>*IOq0b;GEkJfpc!(2Uc#b?%gA|9NPqeb8O=W&asUb zIL9__;2hgHfpcu#0_WJe2F|f{2|TvgrVK8}HbvkZ+vI_BY?B4fu}vB{$2Lje9NWZ! zb8Hg@&aq7xcx-E8w&3z-+_M(UnwX_v*2K&OvnFN=oO7EoaL#Rpz&W?+1LxeP3!HPC zHgL{un!q`?sRQTSrV2c^y3ZF}j&0t+IktHM=h)^BoMW3SaE@)xz&W-#0_WIf51eCb z4XoJm+0&|*LgsJa7HHrWEL?tn!(8b(^!hg{;0w!{nRQz=c*vQ-Nf+r+{JL{v-g7&z zk$DPcoVg2Tjm%XrYh=#Ac^>-2Sa}{64V>qpci=n^iv-T|&?|7BhlK;@c~~fLo`;@+ z$F@e63@*pEMBp6T;(>E)iv`ZH4Gf%P8xT0h)<1BLtzY0ATi?KA+ncKem-ptX1ylD` z3a0KW7fjt(3Y>FWF>uapg}^zt#D9uYw5>Sdt-wJ-Los*%gUZ( z&o5kjKJobDcWhYvtY2Q6J$f~Xo&i2*^lB2}xuRE-2+tk8nnZY>=+z{`^G2^G5ndpA zHHq+o(W^;>^_dEvs|fcDtR@j&D0(%?m64ZnAs_KNKR82#wN}+6!kS+-iLmBZO(LxM zRg(y7e$^zxnqM`EaR2DlB*L0sHHon1S4|?kWb|qh;iaQjlL#*xy_#gBn49IISCa@Y zAHAAHc!lWIB*H63uO<;*DS9=D@XFDvNrYF4UQHsrYV>Ln;ry%u;nm}HHHq*V(W^;> z*Nk3GBD_}gY7*hKqgRs%uM@qRM0nlk)g;2}MXx3i);eXa3U3ftO(MKu^lB2}jnucH zcD!-?t|k$`N%U$G;Z37glL&7Xy_!ULtLW7vnlJrZF#bK5@HTt z;ccT=lL&7Yy_!UL`{>mq!aGE-CK28-dNql#-hc2OM0n@GY7*i6Yy;t4<8?KO@NUto zNrZQgUQHsrNAzkE;XR{QlL!xqUQHsrSM+KU;i1v1NrZ<*uO<;59=)1GctrGS65)~2 zt4V}MMXx3i9v!`!M0oG$)g;3EM6V_h-Zy$RiRNp+=+z{``$w-P5k4S#HHom!>n`ew zQ0tp+)eD~6IL56x=6Dms;d9dlZZpF%FPj^VHL``_Mk&0IEe*%qZ&kRi?W%1*iF51V zdc`Z?Z5nh3Rl0+XF16XV=z_O1x{SGf<#&e$-5o34os2HE+_~t2cQLxucGsc{-p%Mz z>)nekc#opnt<+cbB=0?g>m3-~y#g;9m~n;_Tjr);aMWd3aQ&07p+%Q9JTU1Ful)20 zj=V<(*Ei{osB{+zj&tOp*b4)a?x;#PsQ<+`I=CKsp~&71y8BeR`x;$pwqK>YztN>e z2UNNTHs~Hy=^ku!so5cw?x9APIXtY={fE(I{tvHok1)FI+kaNNs&{=BQHK{B&afkm z9ku?K;qZA>19!CHsNFG!qgMYm9JM*taMa?s!m&?O3F16HxO|44(4c!_rF)Xmr8XxQ zUGOPJml#g1{GQgJdwQk&AEQex|66pyXBb^-duGuEpJjBZ_1Q%iJf`TnGwht;@)>qs z;CzOiTWpz|e1=^RT>qFWYIuIpb!XUxm7jcuT^w9K!!D|H^BMMGb%tG1jhWZZ=MDSj zQe#IRml=+=aCrlFh2fanD-Fk7US&Au?(_4MQRn*?_RZDBuDeu1R42x}B{6M6J8%!a}B;Ou&yQ4?ApM(W)a&pffe^8>OvG7x(@~? z%vcW;jP3meGuC|t6Z^e^b8hzp&bi$kIOld(;GEl?fpczm1kSnL9ysTATi~4At$~%B ztNT;I<=CDKoMU?;aE|Tqz&W!I;zpOf@l5O~ z_a)y~esr#p_m9ElYx)nBZoZ~Z5qlxOFZrp`)$(E9ehx0bFZrcG z_t#4IH=|3QzgN0{7+qrcv(o*mLAT|u&hd5$T%WnrtZSv)&FC_R<5ari8eP`Pc$Mz> zMwfj%L8YsD*YURMrI6Ea>U=%9`2xM#`P^DeFNJh3W3rC)8e>jqIOb&{!wr{l$ixlY zBo(Lg2AQwD^g_by{jQ*7C*dYN@N zZKXS1(akj`zUd3+Vwj=$!S9Tf-|@GQm;pP9~PIO^AG zIO;dM;n+`e7>>B-G#v9hm*GY$fsnZkH=>eagFO! z{Bi~^YGR;{eTy#R^fS6}1Bxzn?q77-bGmjScU+<0I_EN|kOM0_y$rXQ;fQr{!*Rwg z(ZDTPxI72U-BQ8nd;~9Tbh+-_d~Ulg_?E!+^#MQ2RDPE&eyNo#IiFsQYe`)z_JHFU ze>vlmv-2HOf4DEqc!L!}$TH{XXXkk%mLH5AbMvR+$dB<{{gu3tp$+I>r;C*6UPu;+Ug#E3S{oBDK%H5YgA1DE7d5l}~ZRvVO(s zz8g972L0BNLvM)aYfcnCH>iAOZo`V}C$fn1$=q&K*~x_QHZ~mlVH3kq>rD-Z&&>?S ze%Rb_#JPpxSi4&qj{UHe;n)vb8xEh_7>;YvpayQR;i${Dh0A9VYiav7uH$c7+ZkPI zxI@t;o*j)Y+|EUp8t+temtb+E`QD}I^7-JN1rz_Sm9O0j=3aA;g4xHr2ksgETgzBO z1M7N6EPDlBHn5sxNZ{q7SCa@YAHAAHIRASO;T7X`HHq*_(W^;>SB_pyBD_lUY7*g9 zqgRs%=YQ`Zod3Ot@EW0~CJ|mUdNqmgTG6XXgx8K(W^;>H;!ISBD_iTY7*g1qgRs%Zx+3pM0im2Y7*hW(W^;> zw~by+BAoyIhOpKw>sDB6i}fX}^Mo@+So?uJA*^-GS{Ba#enU9_`wijULSIcHynFO& z5@D@n*0ZqIBI{9ji1eaKg!hVGO(LxGh%-r8YoGNmthLYj7uH$Bc_Xazg)>HYRA4oU zu+A;cE@9m_sY!(QiC#@2tm{ZubwyfBuf5hYe(zE6u%f$o!RnZo(QTiJcFfEE1yiTJ z3#Lx{6inXx6-?gy2JRL9(LXS7-scAd)*e@r2wa3*Y!g-$y=Y1}$J+3AZ)*e@r25k&uT#x4_B;+X zc3h7SYvBITz#U$=ZW^`-K92}4$9-gA)e8Qf#xLiT3}eTA*0fjs-Hv;2Ayj-?z7J`9DDyZbN#^Q`NocWb7u#4fw3cpoo@bnt%+s-!i^)9B7;kRNt@4( z_XZai%r*bQf{Fbi69f0rmlke328q74{C-(*-4nhd@Prkfq{6KgzBl2^8{*i{)Pnf0 zGuQop6 z9yPI0?`sI#~F8{;h2Y;47YWO!QB>IzHZ%a{4mBHhT|N)(~JRkR~Z9*x6x&cdx|dj-lCiDf2r|(#t+>6 z#Si!aVEq^U1x@nc<1z~>hYI=F^S;PMcl5zWn(3adSAXnP2Mmu8ECv?mfei z)B6qF2Zm#}c^PPs8&EH_RfZ@o&*}*wG_WK_u2e>~CM?KwmaBe(m_LuQV%@#5FaSs03 zP_vf1H6fkX3psQ#9P{ADqh?K9*9M>8nmEa?r^%04y592l9Ehb`!+7Huu2rr=#=ToV zJJ*F+7Hf!QLK6#dF56(YoZ-m9*}*wG=5IWc1KjwAV-33T;M{oBWfhYHYjEC%x=dhV zAU&>e9qb*x+*}6DRpS9`9Z1H9;&38auA36PfY2rcP|Q(HbUXlDqY@^BTl% zBU5Xx{gWCya+u6;T&vu8kveQyGrjw>P=N?O-_SIkoW#H;v)A zzB@ZOXGg7_9h|e{o@kKC0dBD2n4@V;4sg>Mjyam%aEv#D;qW=5;h2Y+3`ebJHXO0c zVmNB;>{#0!o$>uETAE@U|7Xko)~PV_Py`+X6^aZdC$9Ckc+zDrAoa84|0?5J5E!*Ncy z@yOkAe9b~8(=tc*E8bla{`X%#pk+E&w<7#++v2~+TrZroE^1xc5u#) z>*mEKKe$T_#~dwga)4XHaLmz?hGV>?42RF94aYnzV>oKPtl@}dIm1zFXUCd%ak9>q zH$LH3FdR9oXgG3kHN(z5+nzXAGCo-oD;tiz>BeIo97leu7@x$ss^Q3QHN#=|pjm_D zx4N+-zcma;er`PSa~$hsP2&@8EyGclwHvr~3`g#I*8JGc9-cez&YpFRPinoM;aKy| z4$j$8Yi9@N?AY(mnEc?LH5_xazR3Y@1H&;#?wPKPx1q7a=SGHOjy5(Nwcf;V#ImX3 zsI{|W&AT{RXPX(HaGM*B9JVkVIk=i(_mP<+;@r~sWKC>kIQFI+k9lw$^Y^uh1#W9I z9zM4*9Cm{mxWR@ahiwhVz3Fy_8`aV-+Z&GcvP0pzwyU=NB+ea!>lLqncM7a?!|CoE zT-M#i=u(?qi!OLKqsy4PSAO?s(A~4r9b$B;Jk;n?+hIi)JlyC~>k&m4JhJF^ zYiSqNh`dJy*E=w}dk0=LFyo94oa@~$5Ovu%xc-6hwNK!DrVk8+?tX=1-ueVb-UkHN zH|g$Q=`Ir7ZqeH0j?UjLoE@CAqt?z2&e^ffT1^gcvm1^%I@#m^cZ%Vdqf-sXc&8Z-pQjs+dH9dv zsP%shM=WO;j#@iA*1U_8b#|ul33ry^$l+|mk%OxlcKys85$71=lQnUU;nr@$L$&>pFAOFA#OPH@Ns` z!-m7>BMsc6hGT9YGaPgIxZ#+OCk)5FdD3viIk+MBr;HuBKW#XCKGVQGYdCU$&T!=Z zyy3|G1;dg1i-lu7wS*AoOTpzc}VZH@d{|M&&}dig3D*dCxP>s@o}+bZt|J&S#bSZI?jww zi!S>&pBbN5e)5^|WpMe-_@dIyXU5R#%=oGrvn7Oi8{M#PzBYEOg>MXp&u<&J?+nM> zes4JD@(05)cRw1Aee+Y{y0^58>cn_0_jLZh1%HLE=5=_xYRjL2H3!@~{t;Mn0RBC& z=6YB($!~Xe{=Wsk2G-e4&3*~2`6jlX1MB)RiMkNQhVCR0n=oTdTrjp170lQZ7EJ8j z18dEKCkUK#8$WQ)ZM?uaw{ZjK+{Ov4b?)`mN)BK!|ld zUHHl8?yl#JU3dCQcXy-9+|N*S!Fw8AVwkb=yH|tmOqK31qs!h{=Ysa!z$1(0Cp1iAq=JKkoz1tgh)x72PhOgW%`9hJCZNv7;u-7!IGy zHgL-sjvSXa95r6SaLnC`hGXBXR5<2F`L~uLd|RF0x0Um|OX1aRU4CELi~+Zb;kcHz z+$sm1&mLlP9Qm$leC}Ohx#6n6`-Gdyj7Kc18K1;5ld)s2=Q12Vo5q{p*fE~t$bEGa zC;8pnko%11{=F`Y*WbiKev2E9{G3nXT*JhP&ovF#TIO#p!woImt!6*sb43#;+yf1E z4;qfT+-W%Sb3WmmPd?+k$M_`9#|<~KalSob z;-sFgW^m4DYl(ARlLK~7nDL171;Y{NQ-f9shznZM_Z9doq4v15*0 z4vgnG>hh8qkGXYrBTB7ZU0B!7j&=Q-84tS+On&6>q_JcE+}vWfr7FqCaLYCU@USz?8ge)yA@RP{Vzbovy}_ z_x{Gk=K=dA+4wx5-Gi-vA22k@x?g+_Gp>_AP{sij*#nd8K|+>14^Fa&2w7q{G|3)T zkv%-g9wB6@*^x>1C?QL)9Gzs35wgtxu}StgA+`VDmOH?gXJntxpsj zwLPiAoh&$NcZ%Ss)v1D`Hm3=WTAW@u)`@3JoM#vp_pmc7vS%gPvxO|RIj6{i&lR%7 za9*-|ens|zBzvKdrIr^JS@6X|mfBuYWWkpTS!#V*kp*8~WYr#ag>i8YyUK9f!>%m4 z%uU?Gh8x$z=876#U1Zz%1lovu*fq&c+{3OlF79C?l5E_=I;K7Bx+2@A#bxvV%-bTB zbu&`vk;nCd!{!YY?nc2ew>Jrnxx86$%-w~2AaE zb@ICm`y6n8e5YZb1MnS&eXjSh!^bGYKG)#e4ae8XZ!_%k&6>W|u*a=_KQYcYj}>2I z3r6?Rg7Ni8!Q?W=aLnyt!!ftfhGT9I8IHL?i{f1+1_Zg14-D}u$Q+0pN zxER~BhGT5c7>=<$Z8*mEl;IfLlZIn#PZ*A|J#IM0_LyOhjrHoN;s4pV)b6i+FJ+u@ zm;cmSmOg*p?8JTd>~c>Z?IQa^l08qzvbJ6JAz{kzgsxfFYg%_d+Ys*><3BqLm^9TJ}R={kA*C?`Xt%?v?BXilKouB zQp+!jEci@gqln*MXLu;-z58AMYd^7>v-E3&U-F3Ynx=-30dZFk|aB+kfm29OR|#-S=Q|o zN!II~<8AmuHm*q5lcPrXqpkNXuHo&A&&^$=jE?2})n!dh^V!dS_}J3=VSYRN;luF) zcJ{-E;|1;PhY!aK+1U>tju*DGA3hu}W@kTqI9}Y&e)w?gJw;zRUc#^+J{&J;XFq&6 zUdqmX_;Bp&&krAty}$kN;n>%sA3hxWdi2AGV_!Rd_;BoN#}6NleU19z!?CYXKYTd$ zdjdavIQDhwhY!aq+Sv~uj#sv`A3hwfYG*%uI9|=pe)w>_x}E*-;dl)@`{Bd!ns)ZX zhvQy$_QQwcwe0MN565fU*$*F%*Riu7J{+%WXFq&6UeC^c_;9?wo&E6Pcmq58;luHU zcJ{-EhY!aa+u08vjyJKhA3hv!YG*%uINr?8e)w?geM~Pq-omgSJ{)i9=VsJ` zx3Zu8@ZtQ{cJ{-E<8AEhhY!cw+Sv~uj(4=PA3hxW?;8E^;do~|`{BcJZ#(tUdevxcsIj-_;9?t zo&E6Pc#xg_@Zor{o&E6Pcn>@K;luHscJ{-E;~{qT!-wO&?CggR$9vn^4XC`5X+6*qWX&d~)&7ADaBJ5EASqm5K zz|EHI%r5M(HkBQ?Y-bK(hq+UB;If@Lg&o%ZKv5q!*NgKtm#{-m%w4#)q0VsgBs*UB ze9uCU%$sBzLYDaEOR^n>EHTWVWET*!^u>Znb|E3l+%KGD7ZI}LwrG-FOvus;+3w;( zmNlO1&`HR$c9uvnELmjRgn1+0_F=p88ZsC#Fj$NU{E^q2$hnY`cP z?@*zKzgm8W3q9(+wD8Bid8E)|U3L+AT$7I$de~I?acw(R=utDJ#~dvqVxcZw1xKyB z2@bu}MJ&{;yU?R9Jp@PY${+bDj^6DlY{D%oIO?)og5uisMH>=2Qo@BQWvebLa zA`9M1$l`PBWOtj2?6ygEJ0VLgw=c5b9fT~k-Lc4mcM`JHdgmew?pFd@tQACP1Z6tb+_gOaS*JKsAZd$8cxqYn`rd+n=Y?_-aC@~R2_aV=3C zbMk`FV{Ts&9C036iSsbQ5$EB8Bfob=eq2kA5PDoojuafdq5P4b;+UhOip{oO)-vFZ zHZHEEV}yU|bgZxgcU-b_ys$%UzA3e7OE7kZJ0aOQQP`pWClxN*fjc?bIYro^r<5JI zZ0A&Ahq+UB;If_5gdO_$3sE09*NgLYy0Ak}oKd*8q0Vq;COcmDiM`%%R+2qi$P(W< zN%mYJOAO~F+4F@geQ`mOy->(9_ZKDEi-j!xe@T+PRLIf`+3saRmNlO1aJi6W?Oc&! zxU$G{z4e(P-YFhxnWxB6``^W$gO96(P4xa0der@Dp~w9FE%cZJ@0onfz~8?@4}Z1% znndo@d${n&zS&mju`aI>dR$*85qj8E`EjkCOz2TFrPuIM*%%>W;r0A$1xHU@S2*qi z-4^~vHgn_eovtsK*Y$6x=-!xgZ%Vp17v0#49o@N&TME~<1vdZB{NLKljsM-kZACWP zA^zJ7r*a!r>|pnfWcN;CmwMhMIDFonxb7i`dn(+$iSv75M%Ig5)^iMc7rC0h2a-*<$M2Y?Nqx=ZE;D!H{4v(igTf~HJ(M_qPQb|ixVJ z^rFot1jjY!Nx`vaJta8m^0eSsDmY?!t#ECN9Y4XnZk)eX0Dhxl=go?p zx00RxFL5ue{}z+$^V>xh{7%K^yA_}BC7N`IM>rJl0Scq zh;hjs{?mAWj;Ps}NzXa@@~gzfYi70?uWy-ioAkuj$)>$d794%{i{Oa!SHZDY|0X#0>fZ&& zzVV0Pu=%IpxNiJa;rQJGj6?C!nG^1 z)VOVtU4a1XjNM6!EZ3*@1rz_I$<|~AbMH1q!K~xS4KHE#8^zaDhW%PWEK?d@)vzBv zIv8Hf&VKlCyt>RIea(?)5?0;kDhY!aa+Sv~ujyJNiA3hv!Y-c}wINrq0e)w>_sh$1s;dnDU z`{Bd!=63eOhvO~m?1vA>TiV$VAC9-OvmZVjZ*6Bkd^q05&VKlC9RJSFac_Iv4-_1Gp-lD%8dvDQSj(tC2PjT$)fi>aS`S&&!?#Q>WbvrcQ$jCht88Chx(9JDYvv_c9#U`4Gc#ojYFLbH{O=JC5tzaa`w)<2rZT*W&QQhvT@; z9eXd)Pmbd{cO2Kbj0Gy(f)#F|3b(M}h-Hxqw`k#FZP;_Jyssr+@R=pwD2tEtQ*9O# zKG{1LuW+3z+!BJLhnEx_bGlT8>s;ZME?mqPo8F7DZ@LIsYIAg|&4r?GxF1|b#6ym= z-0Taw^9mSty3^Px^}^SZ{uSH5Q zkiu;uIC9%maP0k?Rk+Oy7uN@RTNvm0gSV{M*-CJ%tF0^CHie7+(c9KI_YdAq*dey< zg&nvZiXHHd!VdR7J5{)yg%9+48`r7Klb`wsJNWNg;dT)|(CcShoR9uOmi=jfaWS6b zZuO}2{<~|EJ+UG?Fv*@?k)8C;)-jwVWa;1ClHGHKEPMR!N%jIE%WD>clI+Domi^)~ z+uLHzt}U{0KSOq~+0FYP>-%l7x3VwnA@s;g>A~$O^oZpPv9}V-zlDp}(mo!#jeVMJ zj=#3&rxr}?dx;phX6&1|hS~SS$9{?P$I$yn#E!rHg(?hs`4d*SExZq~OLB?kK^L!_k5x_hSUdymu76M4u^+c|TUzWL-Zl_D;AZ zgg;_APS_-trG*|n(M@pJtod6`=;2Rska@VB;zh5XhR9Qi4m#Cd{< z6PqUru2K5zB*D>VQ;&9`*6RVATZ%a0W~%7TEI4vddT>gQx|}R>fICHS%!Bd=r~Ean z8a_@HHmTWMqg|qP&DiVqsnqN=VUrwA7aX-#{-{}vJELNAfQXa)W)t}l%TS?54hIX4 zS}Tqm&J=O7{$3Yr4DLALk866z;S5R4saI< zj#?{!aLQl9s^Q~eVUwDzI@%>#*Nna3-b&3b5jM%;Qo&Jc<&T=xxXUUwZxV5m--;qX zVtG*L5zFNje^&^u;bpXOrQnEV-O(=5y4J)pPQ(JYNkwl{!I6W~gHw9U-&G<9xT^(6 z4=R6f${%&vTI4_vepIQ;a1je}UL!c_sr*rw8aJY1^EDAC`As#-rhOvw{+ZD8-~al! zR`}z&+I52S-;??nDLC}{k9LXHJ;-x}>xCXU+#on!FH!!ewc_ZB8--1{n*>K)Zmw{* z2#(zM6S>3fFF5LXtFQ@oo8b7H9i<1S^r*GcgHw7u|JzIC0Jpc`n4{Z84sfFc#~j@u zIQ-oyIBeb}IOgGQ!BOjb1V=3Q3XWPUJ$hcnNuS*(Y{K0yIC6MEaO9wBhTd^vj)?O? zVUwPCNN}u8<&Sx&aob$gWTW*y!5&%TMvGYR_psm^?#9L#!I9s|qg|qPEZCg+fr)cx zRrJml9Qi3dIHkum=Mj+u+@pe{ca=Xl<*(r>+ZZcsQnQOjyF}}naZT+~so6MTlN=rs z9P^<3QL`HNc*W*IB2MxO1Ut?db{O5AI39Ve=`$F-K1ej(y@8!Li<-6&(A- zbAm(fMv*)F#PdRrn!O-6_6g;W+!aUeFAAG*FA0viyjMM}Ep5`6-Tm`BK<~`$}-s-wBSMS9)+tk6J4|IHkvW|3Ksb_o3jJqwhrya6br+Ia1$|ioYL)9yWgx9P{wA z;HdR4f+Ln+1xKxw9&1{~NuT{DY{LC6ICA(yaO9wBhTd;tj)?P5VUwQtOK_}B<&Sw# z9P`)oaFdPJ_n&Zo3xC-BM{wxg)9+@0pD2Vi?(33@>9CKQk9y=BB%G)MYl~dIVdu7TGrTn2%-c zKz8k+~J*RQqg6tegc4^~yzBQe#g`PonuH@6G|M{GyvTo)Udemed!I8(j6|NyT za-2_a)VQPIn2-4d$GTad!YwE`_REC?$G$hWs2Tg^2E)C&t@k1Jb;VJ$jfEby-b8T3 zxo{=UMFdBjiwchX78LohuP-L_*w+^q9CNGuk)Pt2zfQs?+!BJLE=yLpr36RrT}AG2 z-2_KHI}4j|OAC&kP)n6tWD*Q zc~BgGGqAF-3Aak&xF-6pNSv!0=hs*8YC@KMf9K&Xvf$o^eScMUS5J1=D0bVpN*TPi zsJ^S7K2aS0*AzDS`-(Ylb0@9$emM1A*?me(Ctfw795W&bGla z5<{OP+gHfa7rP|cenOVH@1JA`2w8I5HOUSXvh+f>yPJ^ZT9E6oyO3q=3`#K!F0ygI zCtm*!YoDjcQTt(n!^a-NCVB@7J?g%v&}04%7JAHq_sqoNJXGl6ua@88LXUb65&qaW z{kz9=EUe4DgdW$(qlF$eRen5=J67mXGo{BI?JZ)VF8c_MTJI}3^iCJCP_zAn9(CDY zaOAH1k)Pt|-J!xJ+%UmWmjf!?fr2A<|4!~a?{Jq3j(Q#>Y{DHZIQ9pn2dDI?wbFxA zdaPUjF7uop+{J=pjt&tyz#S?$=16_lB>oN)de}T%aLmIIf}_?)3hsZg2##7SJ$hcn zNuM1pY{DHQIC3~vaO9wBhTbh=j)?O(VUwOXUU1Y}`D2dm5&qy#sMtJFaOj;>;Z7DD zd&nt*W8XbhaO|z8367pPy>Rh{W*2Ys*&_UFN2{ z$x@f$#`Q2v4X-Y;ZR{~0@wMe^lAZY4^0mfwGkNTeNV4&@WxmrL+3S+5`-K~6T)gML zz9M@=lD$#LlIKlH_GTeV47ViNTPw1+CE44BEWI)+$=)Gk>6JT^>|H{Z`M*2K-XmmL zxA!JluXnz#A$yp5YEo_M}+u|9AwBs*UBJSX(Xi%Iq+AxnHOC)rnoEHS*AWM31q^u_B* z_6;G++`pM*-x9Lq_I8qeN669(+3veSmNlO1@Sc!m?Yy62_@Kzfy@PnC7}GLOk)!r} zPkrR@@u9GZ-k+u4q8@esNa!(te+xb4z>Y|@9u&t~nN!3` zJyp%%ludf#Ymozb^9p~&xv=1fb3Vb**XnzInWOoI9^5w~PUdexp~oD3EA*Hnl>`11 zM_m>b{+L^($G)WMLSHLA`g#fB554b1e&o;)dd#1iTlCcTG~@64ioYKON1y#DIAZB0 z;v^RJ9mjAhRPaOf$2)TPG#RKgip{n@#bv<#Yh1kN|4;a*PEBK7x%G1^xHiT`J8g@sy6@Y$)F#G5 z-P$EPlL$N1f6~H5J8+XFJCh4L?6JxYT(&cXu*2LbJ8;=fdtry`-S(nBaIP2Y+(Fps zSLS%i!o~W)O=VoP<8{w-OOH&QWTz3b#5Zk{oleO1EAdUAWM>ev^u>%xb|xXq+|Qh3 zXA!dGHfxfdO~}#<+3xH@mNlO1Fo%$3?aY~Cn5)RfwL`oImEMRPwclNE_?TPRL~pRr zqwezvy?$l>_7r-|f%iOHUU$G*9r&|_UTgdW$Zp+XOvDnFj1 z9U%0mnbL!MQv6L7vDCQvL=M>OD7Z$Mzxf46T@J70puRsCe)la zfZkccA8}qJIO04nnsm^cEHQkpth29IqYBpPF0rrg&(=wbHT7-(tca^S*fE{C9GU`GigU zbxL|YQL^(rDfh_LzB_b@WYd4o$jCM$r))Ng;*!Z`7Yy1t+l*Yc8SB!qQkSKYKmVO8 zBm0XS{`g(4>+~y%ottfr|NZCE1@kw+U6M~%qwkkVod3?9(Y0`GdBDz%^7H$nT7Ewd zrM}y=Um?3CAO0}@y9;iYqNi-K2lPmKalVxvvGh!OF_vW$*R{mqr);xZg>#BHmrFKd zoa#GKiF5g+=QgR?3WHpH>kZ zK2{YRHC(O2tzO~Q5FEY5dW!of+?qm!_& zk=r?O-JL8$*<}Cgo%CF#U(xF$IPMku3XWRuA~km8~?X&dkWcpCGNe7Eb;6uWa0KLvebAV zA=|Oo-Osos?B_ooar|3^n zF#E&ghJAlReoDiu8ur6S2g9q`*$*F%SGThtJ{+%MXFq&6UenHg_;B3I&VKlCyq2B) z@ZorEJNx0o@j7<)!-wN_?d*pS$MJg`9j|YX`{Bd!26pzthvN8`;?pAC5P+ zvmZVjZ(?UZd^q0J&VKlCyqTT-@ZorKJNx0o@fLRW!-wN7?d*pS$6MLi4*n6FRckI1RzdQCli+#(n?^Enmj{P-4KYTd${f|A+@d1YY@ZtDCJNx0o z@x<;!rV<=^4;wq79^3(fWBv}Ta0gYmgDc#**G?D<{>~E|wLe7Iggdmt9VR&XbyiU` zVmVys!5vZIj;wG;366O!MM)$82I2@U7+>-FKg(; zijR{j+{qR0lnQsM;E3h43U_+pV!r4tZGGVRg3l0k=$SJs+zGc&SSR#q-zRvMuu1LD zu5jm6xN{2^^Tnq3Oq}2IgdO%Y6*t`ZLXW-s0>Pnoq2RbCTqHQIM;8l@z5kNJ#eCr| zHO}(|UncAj_vM14FRlz24)|JOhn%ky zcHl-9JK*b!Z2J<+(sm}V+l}jDn7O>o@G=44YPhRmuKBkZ?iTRPhPxZ)8hVrA9sz%G zN6Y5`^|Qx(+-Ur=b|%Lg4ClRw>+LAPQR6!*+?|5M-(7`^eZ~FCJ-pV6wx5S< z}&mzb7a5gyZYjONL@cPpRtY~l^7;*b7io1e$vd1e{S`0(S`rC zVCwK$u@(JNzt4*-_zNLRj9(U6@K;4P?yZdy>(_##51zlsg<9Ve^Sq53MVtMKP4)eW z*sOgo-)rLcg6Jumto3h14ns;T-wKX;ey z-YZjk0^FF4Z{t&YCX`2x(vf#Fcy)TsAKa<_Rirv^J#IC-xjrmp_ z{{I#>nVY}F9t1br=$5!+EPebT8w(Ei|6BJuvCK7kLav`bWTP?K|8K2Fe)v;$`p0sJ z_5WAow49r^@n3~&8s`$N{nKx41V=5}3XZtDjGoXYwNM=M(oWdK-xQ-Jv7frV=*EVQRrqPvwuA)wpRYHiwEh$*=aEGRFx$ayVIV)LL=m zFs+D_^*7rnD`%=F+B6-ZE6lpgcf5IMljCpdag z`GZsbsEhh;1A6eMN?kgNScr3e!BI2ikGj;j1u8Z_6LFHC`i=r(`Az5zuukx?pzz1< zIu{ZguZ1iuIP}^*>Xuvg8PCHO5qjjXsNi`1sQi(i;^>LRgiW}`1xH;vRk$StNA5F- z+~H;v9Q9mM*o0e3aQ(f+Hk2Nm(xcW&4^HXvTFW#d2e@ek#~gJQIlwI~IOeE};PAJM z;IP?MaLhwD!BOk(f+Ln5g6rQh4NXdqo>y_wXFY{YxMc-L4$BFS98}HFTR_YaaV{@x z(i1BPjKWkqi*!IAq$f+Ig=6HeLWIYs8vPn;DBXU4*N8yh+`wNaZcNQG`q53Rr=BSU*gWFcb$^7ji^q8aV zgdTIGa=@SBsLQUxA9Jhp*f&&N=xe1%Uk?)g(A!?*M-Dp)J?2l%Eqdy+%<;EF#ovyC zqtA8{9I@;x;v^RJIpA;yRrC%P962aGIHgA~?JRPD>n%9?MfrnM{;125A_vydjpDjX zUHXVvh_kQY&{O`XOO4y5V)F_SC;6$*d?pt48PlwXe!?GnL4U!q-UkQ{z0*alS+~0i zJ#rW*IGztEf8?h)dSbU?Q{DgUZd|8w9r4p3;h#DU7IxtFNOtxVcBsu|r8Y4hxFN~T zUcwIb-@9G^EFi1p(lnF zF4hO`fMmz(-ek`!qAdM$V3IvZ$P(YdN%jySOALo5*~5e^eQ|h_JwnJb_eUn#ql7HE z9i3#45wiVDpJcno3R%{8uETLcmbG)dadA#hD6;DLz`bH$q4pyMhmRA5P4xV;x_$1W zKk9yx&}05?5_-&m_e{2lzgvVJ{%ZN%CiJNH$-*D|<|v`Zx;#bbagDlD=wVak$9q@q z7JAf7>A?*fK4Hy>rN*5qa=_+kf@`=N8>b77x;#|LL4C$H{`f5B*h_H!Sr02&8o5b>h;OIdqbF1{&vs7K^Yo$kD`{!2Yc}MSDksmqm`QfpznLjnR=&8?!?(HGiI8XSa z)>DmfiPqN@YNj}HIA7T0?=Bu3HK9#7^*O7=a)Gc(Eb8-M$-T440h={{>T_1{r#N!I zP{c`o4_9)(d-#NP!Qbj42l87>aO9_K66ZxCPHbK*xQ2&k;}XHqgMW&?#-{qrP`F9Q zy5-h=1~-}DsEdETXVxP>WfM->WNo$){)lBN!Letx6CCq!si+y;WrAZrQ*rk9lx--E zc~BgCd3zBj^;9*3Q#KpTku)`3E^CxA73V-NbCGsPOX@nm0r{)&DgDU%y z`us)eqB!EbTExlT`s^rgfYy5mocjDX{?_k=5fkbWi~5{8^1Hf{AO0pSJ=w?3e)zb? za)|Hacf6N9?uUgU2ZW#Mx40}(4A2#g0 zO}wKGFKd|i9y08E$uxel5p#RpILG~)nWm=K3P$(Uf~oBr1>^gbg30Y=!!g&F498qw zG#qn%!Ens=dBZW+=M2YOpEVqFea3Li^=ZR--=yc>tZ;8txVHtjOR3j-V<+zK9ic}o z?+Ol^?^U?>3m4ZsdLI}UUl09I*unot74G8-_le-}zj>-f+;90Dxyei2{6xQg+H7-t z4}MiJJ@{F{^x)?O(}Q0YOb>ohFzfm2W<35{>ix(dE}q-Ry*_g>pZwWy(PrkN&CGcm z#QjZ*)BBpCKF5vveVg=LhuVFYIA42=Y}3~vBXj;Z_k-Ujo6Zr-4~dJ{_`T9K$oC0G z=KOJN{+Mh=?x)20-oa2d>HnXTp6lT67r}8q|Eu7r^>2d1=I?@|pZ^dXasDZ|{Y!#> z32s>7{uUh1W&RNyHvbh|quBgUa2*S`-&p@){2FuA@c(S=FF5v$rpH{OwH{oX3U}(Y z6W2Sf!o79NgnGgy*s|9q`@ zy`9!%PfDcFcXPV|Z1_a;;w`*?mLE?po^GHOam$ zWQP{nZb|k%AWu z2p@3$iVtvqA&ZXzMHakkk&V|Y_SS*I4%}|V4tRHAhZ+nLcHjmVJK#NvZ2Piz`fDEK zwZCy)46~Q-XV_mW0`F_Mt6{Dc`xx#P@ZN^I8|K=um*Mz&=%aVHem!)E@$vQ0Jq_pU zH+38;IN}~AIAT6Pa9lSI6de2COBQR6h3oJ^LJ#g>!BP7|1joJVp@JjM!vwdJhiBt( z!O=5E6fW)!a7P;F>jr#O#m>uQPQ#01 zNA(9mE!P!wy;S(M&8*HTl5E~KF_C*i`(cj!vAmmni;XE&lciZ@j&x^WNSV0!=I|t zIjK(Pikw(i=T*4#1=rV8v2lUmsKteXBX0HiH`GFL%*#c>CjRcK%!~SL6=Jzq*d!M9 zIWNr1hn0D$`BR^*fcab00 z-^Qa}^wxU{HM?5aB!}UGLr?jmW;O1bip`lsoaEO<Xst|ABS^?R1@aJBtPYk&qupfBy?3O)QOj@+LRagyIrmE4aSF=1Wsr{1Sbe#eS9$xqoN z&L>5j*nCQG4G+)8(}H8~A5hu*)q9cgw|7NvAHh+V-33Q}$|jt$IiMxIragpB;ygfb z`xm_-f@2<@5jBH*R&c!5uHx)p{3(ukP#mvQ>@VV^o~mYW%4Va)@|?&4y#s|m;`H~j z=DmkF4;CEvN$NeA%+aAj5AJypC-dj;?ai?;M=uCH=1Ap$KgCfOf6s6B*WVwqq4apZ zuIfTxE4=|NCY${Iz1bgnFN*xg;UJ;M{HeJ`@7~J3q~6y_T@**0FNrwWTTid-t?GSc z-0SU7xz|(geirjeO3$7=V#1y!7WF<6=4PI;6Y7y4{#2dbOm%uI z<>dD9@pj^T-}1x9JBjo4!kY2-O2l>ObI<(0oAg|g8ork}zb9q%Ncl$Yp)oEs{yAmj zdL})Wq`vPbf6kHn-05!*>1~&}DdiXajVZCj+GBHMGdDip8!EClCfS>dZ0u+1c3Z{H zsAT6(VTZo?Ak`!0to}BC7cYj54~t&g7TEkheezK=H~w#B&L83Q__=a~?8nK!bIjEz zg=^;zmO))FYPLCkpZKgz>Mmqj+t33k7(*!jjd z-`BAFZAJFGB>TOP#pe%27W|`-#ph4S?#~t3Uy|&vLYA8RR%F4y3t9HlKZ-2)Pa!+B z%=KSI7W{XSr60VW1rH#PhwoRAloj<`kukP-eku4^H*Y+w8*x%#~5(mT|EA~ew~}sz5{4egbLm);5m%n zJmA?4v!?vv*$lI9gX6lN)p+&?=VvkO^-!@rcS-Bm7BJZuo8uVU{Dxy}9Sz6W<}(~) za~xx97>=>cYdFWoJie4-Tf$^xY>s1Woeam=7B?JYTg-5b&2fxvQNuB|MGWWInA6u& zY~4*Z#^yN2*3EE?t*hY}+cJh@Y>s1WT@1(AmNq=G*xpL9t!%O}Hpelxl?=z&Rx})A zTfuOQ&2fxvdBZWb0V{;s1Tf=aSZFR#Q zTL(Xtv1>U`ZfiK+yG~(vLxk+SvFXNG9miNVG8|*w&~S`(1H&;^$1&FR4aZp5Gn`{( ze{5f31$RiXZezMJR>v{ctqsRmw=x`K-O_N3)p3k<3&Szi%?*33YHy#?V(Vfou}*Dx z8N<7}ijAoZyDocj+}o!$u7|Pg3DXo=uHSKQpDx+yYI^kD493O1eflID_x5M)UNP?N zGZxvne^RR-Qa=y2I>fp=j&&boIM#i4!`SjF>}ELD-Epk@K*O=_yBhY|Q4dGdYNism z;%2UJvsAce#Cw+U_pIQ^WmaJmZng?HyWqHo_4mR}EcZEt9&yfD;pP$?d;Z)NZXUrg zuk#j;{`RPM_5WCvu+J>RMQWW!@ab{*p!<7@4L(Os)x zeDx}r*w-{1b6dl3%x!hUF}KwW$J|ym9CKU6aLjFG!!fs&49DD7H0-&lx^HS+jBOLc zF}95j$JjP99An$iaExsO!!fq?4aeBlGaO@E*YLz*+s?Qc+qQ;dY}*))v2ATQ#9CO>=@WkrAmvJ$+A%nYYfMEIK+I#c{tc`oQH!9$9Xu=aGZw&499sGW;o8n zP{VN^_BTAS9y!*y7~3(1V{AtojV8JS)cy2=srzY$V{WG!j=7y;IOcY;;h5V=hGT9g8jiW0U^wP>yy2MJafT;W z_lu2-v0Y?1#&)6M7~2JgV{GRejER|t;zzfy4AyI)n|uC8#y3&(SC4+ERm z7#IJx=vu?R$H0#ec6qJyI>Aw!krnRx3U@<=yHRk&a#Mx7xo|Om^g7wx$NXf)gaCa9j=8sM9shI!0hU5I-BkZ!5`FlxyJ;t>$ zz#p=4pU@lDLNy(HxBt*uZ@F^TiNlA&x;mq(JG>)`KE!BMlZg;V?KLv~-{qhrgdX!N$tO?a9>xrZweRl#isX6oZoMS9r{_t&GqLy zp+~QOU*UczT#OC9AC2?az&};&{46-GMZZ+IUkexgqxYL}?jQWSutRKr2s?0p7CYd- zgdKY0?+W*i@PXdH#&z`jCrz@~R%9njvNu&^KeGFp7{jeXmUTXPvO7x1?pku2BFWw*WQP{n_DS|$A(qi9Sh#5fNAER;yHM+XqONn@KOu+B zX@$*(i`$q^a6?+Crs)MoEHhNN87tgO6>jDVH%o<^RdCGTY=w*K0d98Ve4T;k5O&CS z&I&hIg_~P&9ZP-Zsc`cOj`#mI1jpaI&L=qX>sUBdi}{U{<{gA)Z>?XyfCfzA~^LP!9Ff-WBKd+ z|JHiMqTa8_{n~QlCe$N8{HZ#1Pj%`ca$?W!S>cuy9M310t8mL#xD^D)e&g@q%XPxv zib4-=r3$yQ;PAIfg8>Yq>D3LA?nXryenX)a{sVp(%rh~!fz#X*(bJ1y4w|9_-%zQ_ifuJ-5rZA{0>Ri=Z`w>WLy`+$o4ku zuSwu{=b}sB`s)t^?eNjpxE^-KR-Yov{<*Ai$nH`&>~uAby7V`$TafLSWc@XnL3a4q z)wrH^9@vbUng%3Uf0(^+pmDKxcPp~Q=d#4UyRgIcW{`35`ZHL_Qm-A;^=GHb^=FS{ zch6!s#>f7w-Zw~|iX)yO!Y0?4tMBy+w7!nOsrN1Qad8_vU+4d~)*}}6{wivGQsw%C zKl;NxQ2)JBo%R+vk?%f&Bes1D7k$9(SA2l?uhTSixcQxWdK$fIIrS7TfGu#|znk zEh0_F6j`p>CluM3$Nok3M8PpvCkc*TIk|8#R$iw*;kp*%_&7z_A+}QmN3Bk)aHk87 z`kztZ&J-MR`g=6;z5>VlHDV2^(^=}|MK2dDG~mbv!#w&eWac%Mtm zkKQ{+}_5R9|H8pV>6XIo>Sb3k`1`@CAmcfj@k{VR{9; zPmn*)c&_>W@VSO3wqIUW@?cN;QRG2Ael>1GH&w<6vm5hp9P_xwaLi-4;h4wOhGQO% zV;)x-j(J>Z*z<@r#-7I=W4qa8V{DFNY&RK>vE67GTkhos!!b6;F}CXs$Jj<1&apB7 zzoppjGT9iL;~3kWhGT4Z7>==xG8|)b9AmrPaE$FX!#OtQ^!F6oLna$za~xxP&~S|H z0mCu2`whp~9LL!1GaO^P*Km%Fed3Q4+c=Yru{n;hjWrx&d(?1@?GeKyaE$GJ!#OstQ-7y9 z{n}(>Y>s1WUm1?EeQ7wx_J!dXo8uVU=Z0f!pBeVp;=Kj0{fsNmPZjrAg`4b-32m}3 zUEX|MYV8m1iVAn7;E3~;N-VEdxIrRL=H)8kk2+s1IC2;+IIh9h2#&cLAvpHSYX!%C zb)Deo`;mgf=JkRjmKze+)2FkH{63UrAncqufBlD%Ie(l!yfOLn*M=DCJtNfmrljXO z%+bw>>xw};XPezBoKw{LmSod8p4Z-5;cgQgwZ2_&1KqHV6(0BhTdx(kQ9_Tp+#xvh z?yPWk369>qTX4Lta*yD6&G25qF-P|mPThy!Z(MvW`GJb;gGu%wAxmvW7g_MbLUy2+ z-Nu+A3x1>``)HCKD`ctVxFQRFOvqB($BQiZ2_Z|ZpDeQAr;04EzsJ{hpEjjM)}IM(pqsMsd4>BzaMbQg z!BMNP1V?SY796$srf{qik8hy=XX9Jr;vV*0MfUq7`-70BHa`|w@J~W^U`q&1KNnf> zFBREelk9IomRkN^WWj$3S!(-dkp=%HWU2MvMHc)|kyU%xzsAKqtm(2R!$j(Xh_}cm~5h*F5i;-muR#csj#A57ca0!#>}{HjQDgy?P(p{Kk2#`07|N zy7Lu`uSUVdKCj`J+dPJ2ZgU%sxy@xb<~FC{nA;qNV{Wq>j=9ZdIOaC1Vb4v~y_0b< zw#5y{*cLM!V_Vd4jBOFaF}8&b$JiD!9AjJ1aExsM!yX&!%u&PtvoVFeKJI%dTrB9YBvfz7#EHQLWb|0w7E}djY3t862)V8KPH}E4u zmbKZX*aeRhvaHo*iY)kvBFjCUSCHKN{f=IH#dS41nWJ9a1c%M;6|RTi$iJuH$ah)6 zk>7HHV+}7~xY&opxq@-Aw^ppku9Rd~7P8c4l_CpXRmc*H#;hK79)!p{apm$`}8^o@<{VY1Y4qasV6 z#%uZ}$&T-9c}Uy}8D=W9fx)PBlz zJyGAui){a5hdy!>yzVnVaOmwSIR2(}px~IJ-2_LScdu}RD%{`-w@2YvQ=T9G_B1YD zZ-xk2;`aCM#=TL!R~Om6L_EFSv5mbe+&&d<-wL;%;D}}a3OBTHu{MKR!ff*Q>UzFl z-kTd`@!{`>jkOsje6pSosBi~XxPt_@cZv02!7;~&RJcPc++l@_`C`+1G4{>jLYCV2 zdwyeW)O&o{JB|?X{4dA5J=xaxCDc^CH<)|!BZW=4qY~%W5QcgWEdGv8dalDhbxh*? z+QY~;{W`?RoY#y2Hx&=qxlw*T*8H^(j~`bs{E3D;B;Il4c`oE)uJ1VV{A>?5^N#)S z;r5qk`LX%`lWw=yRs%l4c*l+M*ztz*Ja#OcdfzYee^N#70VlN z;V%)otli6!?iEEB{&JzqeZZAT_v)exe^t`;`J;}*jf?k`BMirT%4>=)`&hiEyw13I zPkC*TZBzO&-cycDcH%wd4aUWL%IlMCyr;a$IDekZT;7;u-8Fm3&Bn#vy`{(!pUYB* zTLs5`%x!|>UgdVdHQca`Q5Eiv3U_CPyQ{+8UE%H#9Cf-^aQyyQy{8jEhKFe5A;A&re|NY<>+5F2AF?r8 z=rx*|rlyB0+?c|}Jkay`jd?t3I9|&h5q8OOtl(HD<0{-^74GrE#aOZVgmE#}rwqqy z!;`}9pi-Zw1-DP(o~dxpR=DRX-1CAXmKQ4Ai-n8%qu2CA>-=9bSyQQ_XKaBmeZ=8sM9sZM3i{Pd3DIR9@8yY#BR4?3?a{=Vd0p*PIM zZTxee|KEB&?&lBLcu(k2@An1Aby?XYmPM{|%dP#9!w14r&$jeNeKQ(jX--GjWMfR5@`)iSn*VDme4u2CI>+pBM5zikL?#~MMSB3k#!u?a> z{uLZ+>A%9on!+_b)v~XE+ZfJkl6>1%xONq862Y;yCarLj366Wi$py!JP9Zq*YhO53 ziw?&59*gXh6+2T^>`a~POe193SEntq;OUBN>{H^IUf6+~q1XY>DC{(RYHiF^;bs;b z_wln-xLGUQY!&Xln?0h|*8u$8U*R4Q9DC2~!XMlm6>d(!(fg~}{m8`PoJ;7D!`u~a z9>Fnx^A=9^X~Vd9Z!n**!~V6Kh@0zgN8tl*{^VnUijM^ax4)Om#zMje+``GnBEk-J zTvTw}w=5=nz%8D9bP{%`)e^!Ezi(f%nQLky_JO6Yaek1^K9{Qa>@0lJtN3JZSh^yc zeRdH(xhGpj*r5j9nz`|R%iXmi+ugX>hdqjH+`D*}+1KvqiGArI4l0TNPRG)0ROaRJgu^Bi~&L7h{9#XIz}u{=yD<45)Cs3LoeVG%n`9 zTe9Q7>*4o3`THdY8Rv(O-IJc%$H$0qE;0W8$@3$}!^ApjxVVi$$)-QXS{W=j_VGOg z$J*R8asJveV@TqRXUQ(vVAr}W5oe-R7ZP{Gk7${(EaN6pmxqv_XS!XLRGAUNVYP;lref7C^B zTw@OsHp$^&!I8V-Saa{+=p8kF%{hYmpu&A~*MxfHaEOS7+z(A$TvvxBF0QN0xp!xzLNxp9`9j__XON8oGkR1qf-ROJScz6pW;|6 zrxu%WABH>4xOfdbUHGR?X9zoRXC^ym2|Lur-;d<$GsXjVcCvGhutWXNEnKt%cV4n{ zzOchRwz314?OY)2Fn7uhT()zeu*35ie;-(`51i}8`MOBhp(idbT&xe=CCQH0J^Q3b zE={tR30dO1Jjq@mWQpO*Bzu*Rr7x~dvcrWebAL^e9U)}N?b;-Josgv$vfYtFmNlO1 zaJ`Ua?c9)JxUtB_y@Pl^DZLRnYOmgJg^!zrP4xV|Zuxpi-ES6p%%8vSE$cA{-ZPoQ zpT7q#bNH*}=kJHhder+C;g5aO-y4_pSeLg7J+ArwKDn%iO_d+7YxsNSvK}>4dd$&n zA{Od$yWptxD8Zq}`_N*)P_sLP9(B1>aOAH1k)Pt|-MfTMxVr^MUGAxH_X>{O{XJuO z9^iQYSgZ^6yieGKyI*kZ4@wVC=}~K?2dDG~dRjL8y>>Z2INo;`^J9)45IMj-C^+Uw z{Vg>99uj)k94$EJ;bFm1>oI~OmPZ6ft(6`#tM!c#tDubRL#(vTKrCs zI3E)>>50b$N3E4V=4ck-5AKPI%_jwi-cuFsX~D6FJR>;v-Dd^I-uj&2=!xeG7w>n7 z^9AGL-)Fp7k$ox2zAR*^_bWvf{Hl;8hS!qa*DJDbB-uBGEVX>A$b#P%vefpSA`5<3 z$WrU~iY)m3BFo>ycs1ul9{kd`7g5A%OZ2Viwg={VK46mbCuFJ7 z_eu7LitLX`_9r1r&3;a@zX(}+<<}(pn~-Jxe^0W12wB$cpGns1ov&-i{v|l}=)VQW zUb~Fg``Du=yLm!?TuT(kocMb#bDYfW^nxSKe=2eQD>&l(PjKY7yvUDhNz>C4#>us$ zjo|1F<&XRn#~igSHd*WL9j={mo&13^CK3Lr)1<-<++@klp~w7f zBlMU9@0q+8<8M2mhre2WI|x1My@2q?zPXdoV_hyN^titE7JArJ`Ejl7EA*(D((CA@ zvayherP0hZH7zVS^cE4^kfOJ!;K*Sy!O>@nSGZ0UZix!FWQAL*aBBbQY+U@^^HbvQ zuE}ZXijOW8ZkY<#wZe519ImSb-+*8=$_h-vixRDQe6X@T&sqa0d zN0$>e>5Jtn+zJ(L#lppWvFSY%=XWJxhu=l1xY<`$PI^7d8t@bEyJSy|zZc7#Klfo| z&L2OdkgF7blOb&9M)~=~+b4W4XNMridoy{+F}nVk<1vtS_V-@KxxCzA87Edax8r&- z#=VN(suq9y^0S}ZPKSi2PPmb9j8VNGl;8KQR%$hE(f9d2u>8D! ziqF2nC-)%u~&ITql~)&bt*^_EM*~XWCuZ zVQvRCbK|dTgDbLo7}qKI+_T6=pWJT^5q4N#dnG%23t4KiPm*3waeYm`1ZT5v-OcTC~h7wp#~*3qfP#p}l@hU4|)WW(|LagyP9{W#HZ zyndWuI9@-FHyp1Y#~F^-k7Et{9>ldcUMnv$&N1_RVZqeuf`YkLonJ85zVi&n+|D%| zb34ay%>?pGTZW4p?5jO|LpF}5oV$Jj16 z9Amr8aE$Fz!!fo?49D0mHaxM|ZZs~&c7x#<+x3QHY$FZF*se1iW4qRHjBSMB7~3_5 zV{F3>Ppl{IH7@qVJq6PfcNa`g+*L3=ai`&!+Z~2uZlesx+-^4T6#e(s> zLc#RN@&(f)%NdUI@SORI^YE*kOx-^!n7V&h zFm?aHaLny}!!ft_49DEwH5_w$$8gN;ZNo9Qw+zSJ-ZUI@d&BU=>i(^9F}80E$Jo9$ z9Ao>+aE$Fs!!fom49D0$HymU8%y5kDQ^OvcS{wd(BynwgSCRG4y@;~vb1r(jD;wVz zJ8?hQvxREO$ZqmZVe6i`h82e^2Axmw>rP`>^)!@0{3dZm3 z1+$;eTrm6jOorp0IiumYXUL(5PGdOknNu5%d*)PzbB{7d za|(_+%vIs$u5j~IxaV%2Fc$nhU*TR5968J@{J}LU+mSuf+P3t6|P5x>nXS)z7%XMTj7=y+~5|qrsW04+^!%v>at>mTdBgWT;W!!aI03h z)dWWmu3k8`Ue++q_i$v_6n4nBSA|=v!mTYhu0!ioxOD|bETeAr$!mST0Cz`)oAKTW z_1LS{6aL`VuW%a(j#wV8#4@(R9ao8SG4Z!l#JRZOnA;6SoLtv8s&L9CuW4>9^x!rT z9C2=1;g%OU?C;ZNV}%O0S+O}~i)i!zjq>y6&0JFxefQO6Ed=)#TNuA#k@eFFBU^uG zq2tX0f4t$%13u0$b@qpkHOyWI=Dx=9F~;ND@zI8JKf`Y+YBkKoZERKHwiX;VJ>Tq) z`!r>f`fnuaPoB$FIAwD`55dMZA{P3D&q9gMweYuO#h>a~_Vg|ly=_G-%q#vnw1nOK zf1~`Y`s@F>~8o6O>KXuc`~z73GiXLqFk<-1`fT9Mt^L zqXUE<*Wz6(+`z)cdnELBGpNy zppYkt*x;rSe{0CTce2nUC-qq*)Z(C&gNFh4&tr~SKUZQdP7(g#PEDM@7RETN_>1$~ zDCAaSTYjL2TJE)5`zN;3k`LD;-_sKp=M4W`!{eM~F3wrz+#j)=k^DJF4rf-lvjo?0 zWgBM;j%(vNf+Lo51=nb{*wl2M;MhA(s`P>ST#x_xtLXXXKji*lj!qLC`SF<&%#mAx zQ#QFTI8)dpPXFwO91CmbY{AjX=Zl)bT_8B_rBs~U6Df{)P#pX1g(6PsscHtNY&N_c zHZBx7pr_uqPn_3?Scua<2O`&+`(5>U2F#IvHb&NiyGX>z{Q2i(WIg8SVxh+zsT}a9 zIO^h`pOO7Bw@Q!uT~!zQTItc({#hB>A9|OF{K(-lp~w8Gxkc}jCnoOi(waZPHQbbq z%LPX)>hlwbMSYe7+%@QO|UUg8`l`KkBq6N~zs0M^4b!XJCV2*I)5uN546&x>AS z-Cig3$YG@5cuuMOk)Pt|iR+8ac#i~kgK_csz>UH`b-GE|fx9``xkcEaHt&_%#CYIt zO?GY*cBudDg^PCJMkPCU2s`Yt$_`w%bEmMw+$lS7+0I?U4%fT4M19~~FV5H9!VW!g zPvK&H;OP4xWp7P3w1 z{es?3@1i3|WtL`K-|68s(qWkoB;s z@*C>nHvDrNvK}>4dT^87JfR-3)VSwF4%mENa1D23;|0M{m;Wj`sLx5jUx%mNvc8X) zQgHNjJHe5=vI(bb(({uEo5VT2;ON0A1jqf(i=t+5FA0wIpyFg-Q5^H2IM&fLB2Mb5 zY6holHhkJ_yex7+ZwBFyIOh-?an2;TK`ri@)ceDkqgjL=+$$nZ=5IEk#~i&X^q3=+ z1O608UFH)0m|LaCo~7zSUn@QOx*`0b_nOF$9A*@H%%7TD^mu>&l4Z_1mb0JQJ=@f@ zV8Bg=mohx0JT{5pRRjN@#kX3({~BIB;C~FS5%AxJ*9`bC!@UCj)9_jW|6zFTfPXi< zUckQ@UO(Vp4Q~+eFNQY^_-Df#1^koYjRXGC@TLL(U^xDM{(Hll2mU+5TL%1=;cWu` z*6_9ge`C0Jz+W5g6Yy7t`v&}_VP0$Zhrck)JUaf|Fum#cGsEa*S0qp#d;lTlaXn2o+KQPQS(;t4{FumdUJ;U^%<97|mzfX9_Fl*QOw+*w-9ee$F zJ=^h{hR661Wwf#Wvw!~^Zy5jF|4>uY>xMrJxNWm=>+4s?!cSVbB@5tG?naUCpZIIe>Q3~vzPncwh+ z0e3VU*TH;-<2q;qsZxedp4(B5!d2Xh&Y>tIg9aUINIxKD_0cEfQU z%w{;QgINv7buf$JxDIAE9M{22hT}Sz(QsS`GZ>ERV0y!YLp;+N-Xq{?4aapbjp4Ws zrZybc!BmFhI+)UMTn8Nt$93TKi|b$t!()Q~B@NrZ|BcBF#C0&4;Sc}E+8*94+B@T(=6>y~6dVa6JXLM~QP;!O=@gKJ7xS>jJk_h3j15)Muj* z%W@(XxaAAiuEgypxW&z1r||d+CfhmiD;C{0L8eiDUP*A|c)$1@7xY#ZdT^^$xK#@m z^FYt%H|DXrVb6oHny||{SwnD*qPV8uxG(D^IP}&kT#S|b+_jC1v94>_V})O*qPt$w z-5}|%FLc?@HcYx37v2BI+*?4|RRn9elLQM8+;s>JLE^y)E+GVhdlH-&@!;<6?(PuW z3GVLh?(VN@ran^B`O|mZx7K^>_F5-r{_|H=S5;S6?=uX0!*3{b@oA$%w{PgeZ(Pyc zq|n{0(A`w%(gQaybhiv$_$`Dk{dlWFcbm|K-&*MM-F(|Zcl*$V->%U0^&^iv7}v=# zvi%KrHq5;JLbtVj=A)av$m5R2bvKNy0U^u&TQ14&RM_ca9C_KrxUNZd=R(&1zn6jb z^0Aw7J^Z~1nV*HMe>^z2-Hpq&yGO|4pNrx9zzVnRQ!d%EKccst;Fy1qut{DAS2$&p zb9{)kB31bH(c*1oAlzb!cW$T`@G~D#b&?Ce5z*HBRf^}_7c7fZeGF$ z=40>s$Bn`1^=-LFu%mn<-o3^A^aa*a&A(5RURhsye(b26jup8fZ~F@021fk*367Zd zuW$zlj#@iVaMbTXf@4n{EI52o@ePg`4iS3n31yS|Re7PWsrl$nhl=@#{V>51gIYgo z(Y>=ujT8l%L_YQ z%nos1XVQ(6jd)OO=Z0=#)ONq64tC%y_j^Ed5 zQv2p@p+_9=2#)#Qt#I!Njk9tDu&)8vh@TtwH2{BR*w=cvy?lIX*w-5TiD5rW$l1q+{mdF{FCQNn_T0Dg zHygf&$o_7eW9IrTFuK16X0BfXI*M>bdmG{=qw)Cx);q06LQak(hui47J{bM-$_P62e+h2yWZ+{xjzWrg?ePh2m zYV?0>+-iSA>AjTk&13#kOIhl?jrp1T?qN4hT6VHR_DCVi+D{&`;A4a=K1@;AJ-#A4 zWg&Z#kR1|%-#e~3Ht?xJmc2Pu*ae>$i-(+iIHW)R%qW|Ivw2AAs)KW8#7*VfDx*;xwNS%oaQnJr|& zvkO^#n4_>eXGM0dLUwK;J0#+sCuG6%3R!YHU&w;z7qaBKF=W9DgseK#+Z&h9^o0z2 z4`QALLzlJ5XZj+>bvIdZxNyi)r}=xyqJtb+A8rLE8{H;SpcBw*kX(3CT%M`N93R!&USjcv&$aXGdy9n7K5l`1bwwsV;4Z9bz zJ%lW^vRom%ypUzzu29H&zRNSBG5o!wI8Ua#wON*0am1OuV%W(XbFNh3R<3ZXRJc_u zT+iUBfA?!R2_d={1wrw~w$x%__f# zN4_@-z3BtH&AsC1xh$vN-{>uK1m=2o_Q3cwTi_v~zp>eNedd}qc>J2haK8SX*|7I_ z_?ZlMG0c238qU|4GZ@a-nA02fp3i*K8P3<3?F{E@%xMjmYfUVDBQJTM_j7R@n*_(1 z<)6`0H#IKz_00`?PT@DJ=x$NyZdK@RDRfzztqa|4Ll=Gi9lv$12LyL0Hm?_}85C+C@SW0yjAx6p;(Rp@e#>|W>&3|;s=3SD16@;J!2 zd`&XMu-7c}4h~)RY`!Mh(>SkBYz+-rYAj!q3@hyTJ{V#zA0v#**CfLW*?dhh$~fNx ztmVi;*6mUAqm9e8J0@iD&t>s_tl-$UdkKyn`sj^rMfcaoD%?_cPppS8dyDzt_Nj3D z3JyPCt@!dAy)^Kk;4co0@0SQ4#)jjD5A)js7yh7+ThRtn69;u-%1V{eW zee}o`InLKPWj_8l@#D_`;m@du<4O@DXUei7H~6}i;E3@mVH0j=;S2e`x}v9Ip#Buc zeAk4{wh^ztwU3|u{xrSQ=l>V({m7bKD{=vMU4^?|aP*AbMeebAgV2M!vBKRXIOZEy z;cgDDb;>__w-}eNKW`Ow@a?t=cYB4qLvYMLLgb0>l6ML{YVEEHcXx%mr@}pSvnQzK zbu#lkT;U!O967sJ%m;U0g}YyHL@o zxWCwxd|>a$-g;Z}hdT#%S77?fT_)qd3mhSA|X1@UVw{)|PzpzsAh{k|+^uK9kU z?k~a``MqjNJ-GLTpRAwXx2Du%joue}tdWX=`4mT99v1Vlwn}fLyJth?g}PRH)V1Hc zri>504@7*#z&(U=U9)~_ZP8QrhasQ;ik?MY6o;Q53P0(sr#(Dzd~oWX34>hR#-Zc< z-O2E91Bpc9m}o!xlBlb&qpD2S+^F&Tqz12QK?-VfUBNo!QA4 z$p3C!CqA_IX`XAoUgM*-hI<0+eJaB_#wiWw7^g5iAjwZ|n10|NPi8p#-o`Nd!ui&Q z>79;S8P4n6U>M(>|JVE<;y=Xr$M8^qz5i`E=lCzfd9D96oa^Qf!{d|w)Ha{*D{3Lf z+OF_r+K`oJ+ooT^-!a?TbBS@vVj-(tb3ec!>j z9K(`^^V%(8IOk{aur;|m6N6e=$~b;@s-xlDPnR{EZ7*Xu$G^1U9B*F7PR8Xv*g0g| zxQdNzx0`V}H(d>Lmb+}1&?Q%^`G+yO2iH1q=6V>n27>m^`9NN?te*kvC0RcMdL`^< zz~%{a1~9Kv~l}Ut0%q54X>53_wekC_waQRzpCMN6JEt| zUMKJ2+4q$VZT_gaQ?oZhqZTx%Nco#OJI-6!ES3~!RK_v{>}_v{?!>V`eeN%ic!#@@5DZ{D-} zq`ACj=Nx;_&b9bI_3WItg-kc+rI+D6xA*KmY3}t6=eX80yk+A5r=Fd2M$c{?`q{4y z%x=zS*~4?byoYC9@8NkZdmGOA`k#7ut|RZ^x#xKg&$jy*&iUBTaIWX<=SIeDlyX(} z@NBoQ$>v;bY&#j^ z_H|~h9s53DFF5vkqvjlMVc1_j9Q%G`PdfJAK!0%D&#=FIIPP!n{_^2?fW7<6hvOaX z-CsT&`?|CCj(0ZfFCUI~v3Gy@aJ;L%`^$%8uPN%vv429nqO0#}7@r*P61G}7u)#X`&wB(%JlW1* z5)t&N&`tIIs1dhI8EvH9Q_= zd*9!1K0|Y?2N{=rIWT0W3H>a4jBz1Hp=j~L(_~GIU8qRZ{Vt8n>bF$$a*GY!yZ!UJC;k=%XsLwNu%Q-teWZ927 ze`f}#_WilW}#gy5wp#|1id7!L<&Yxyy}P13`P|d?2q`*3SU0 zom|DwfL=-0&w$Mn<_ut7r^vgX0h=Y+vgh|q+?6KVC*dm$d!Fzy%X-hxK3`>c?WF5H zzi+~2&)+I>WzWxcyytI_WWDESyY&3&vrPQ#JwN;GJv{q-z2Vi9p7-#z680XReeoW? zPU2TJyl%qR8P4nEJv{qF!m{56LABz(2u9Gmy>ymsEh^SXNv&wD(_c7t(w?y_fR z*;`FE=iwH^IqsVc=eXB0oa6MKo#z^7m}__U(0g_t{~^XrhBpD&yZ7uI=OKo3oHrWw zI49M!^BQ~4&c1oi?vv*7o}F{-Jv-Op|J1W{-n?h$y!0}h=k}i6C(V7k;T+d(hWXCm zV*gXm&N-uJw+{Q+uRF|c&S%-fbH2QXXI<~%c`ffWob&ZR_3&Ir-otax^B$gU-(@)G zV?)Eap0l5K8@Ex!>ThKa&vx%M*_^9;4D;R2b;}-pgMS#~K9hAHS_M8N-u>G8e7_klxPGo|Jp?wasphPxRa zZZ98?8t!iI;71JiFwFPLhYc@hnCq>F40|rx`J0W7_Coe`;~X>hYk|>yH869%5|~(C zHk@O7$#9PCMZ-C^7Yygvo;RFhd(Lo z-Zh+kd&hA0?QO%^x3>&u-`+HweS5?3qYc<(l|4Fg5W@U~1y$z|_P~hI4E`8qTr(U^vJ2z2O|&cZPFp z-x|)bePcMs_O;<0+gFAsmG@T9xAd*SaQ5wA)6Ks9V>tWvx8dyDUxu@9e;Urd{b4x! z_Pb&CjlSfl(f_g0+I*Q5*QUZv793}YTVbvIe$~q+ywWqyPWbPTI&`(;Cj}Fpc574pSS>>(JJ4UWcg+ z=XIFUa9)Qg3{R>?W-~7PHml+6+bo8&Z!;UtzRhGf`!=KD?Ar{6vv1QI&c01&*nLxZ zU(mRY#?r6b2PW?e1Rm}tY%~TYuk#zuvCU^V$2PCw9NRpGb8K@P&aus9IL9`p;T+o> zhI4GQ8=h3&I~bRJTheg$Z3)BKx5W)--xf2ReOuIU_H7Zv*|&uaXWtewJgL3Wz1SN| zRb;yrva1<(yLr!#^^a^U9dZ4O|s#kP2u0z=Hw=TjC z*GYSbn05)>u3`?jZp9qkD|7S^+^CpiIWY&^^2Ho02s>Q=tSGoKF~>?`4!D(zIaU#N z$knRC4*R!f6F1>^UAsqK^ZFoL&beA;&eg@7!y+!`q|Vl;$d+@iDdyxlcP(Lu9IVsC zP55sGYgc5~H7?iTdLf(l2)|=nU)W)7do^(re&4!5MYgwbdCm<(HqSXa^4CY$VSjB@ z*x6Xfl9RrL>?T5%ST-$WHxshd|K^447DARf-m;M0O2`u1)`jdgLY7)B?QSb%slPH0 z+X-3v$M%H}I|x~F)~}H5BYGCH)J9%!;_e@FWRATwKycU`E#~7K-cjgL(@GC+C!vR* zwPyqjv+F=VgZLewqyGMnjoq4TPWZj`u7L+N(+xWZ#`j%>51ePa3$8Kj>>)Vx_6+}W zj+t+u(8K0FBFC)F7{ReN>Yn=Kk^J*}Wp}^K)&C|p_@nM;&Ym44VjLKG8z6Fnuj*ds z#5h>kgj4rAC*MOVdMXBLQ*q2URM>1E@r{gi;&;FP9fbe>vcKtdG>n~n1Jh^r2~0lr z4$PUfUtq3>_X0G6TxsV-VTZL-cHl}oCkZ=T%bX+f1Lt~qy-pT(sEJd8%lUyjwXoxPFV}<`IjxXA zUC83!8HMbbLKYv+DrC659oZ?1e&>S}5&aBxKp+WgadT zvh1Bp3Lh>F+1xwu*Y6o$_Be8{?jO$_mkFEb`Mvu~o80ffpSBuZ@?_cVX?<>T7^i99lf2qg5yi(|K2Hz(1u&LtXJB>dBpq!7KDLvNcD&Y%x zxms}K`WnHZ_n`2FoLwvQ$jfztBX%_(@hNVwr^Cke!Y14ef+H_CR=Ar4N9=wN{Bj-O zxF3GrW8`z3unBjw;OGxZ4^HWkYo!OL^w>xK?0_;pIG!Vr<716(5i!8sDmd0i-P@J< zZWDUgyj^gt!ySSn*LMmIU+xkdxmJ4Ayz-MeyIa_VyGL-uaIfHqLFEj+55@Z7=Y7H^ zHF3Y-$hDe}HS*^HOsZZUsMs7YIBY&x;T{qkJ>+4*(RUva9KH2X!BG>B1(&aP@bhuw z{JRkNiHhu#h3r#8mV7@QvfyWgEOS0v*nO@d`+On$f{-PbFNQ4mB_T_0Uk+LDD?*lB zzZ$aO*Fu)>@}5uPe%-iEhLL^Kuz&Yso;O05wdt1D;%(!)CtGiYEWOLWfnF*S^9nSm72D9QC=d;C74{Ruy@H>sjGet8nUBAo#L~@C9zs z;J9{n7vWa2`Erk4%w+xA8h&wMx6uvQSVC~b__la%Nn?{%L&K6n53WOnTPnC52YSA~ zIgVuvdmN0VLpS&Q_CZc@Y4Z=%?Xp6ad~_`486(yUJDtK#?rr?{#>Fo7Ia~fM5p#5| z%-^M$zpI#o-lKe@M|Kl)z;!R?=uw$tIlmO9?5klk9y65BR~?6yLdS}pBvCuFI= zG7sAeS^CEgg%ACNEII36$ZlKdr_@GXZ{i*hb7XE<)clTuTRym*1jiZnhxi^1_h*Hh z<9;{ZvbPV6`F0lb@iz;*2=0Q=+f{Jb+)Z%U++A?k+(U5WWuV|1!_RPL zhnuRxt#tpydeq2ZF(2HJ3O7`6_%eINmpLlj>B3L;$ev<8xM6~$c83d&br>NyVjo%I zMpd}c74F_!C%q2$Rk-^FH_T1h7$fF`8(ZP_5*)GbUE%hraQh05HQKMj?H?TfkG;p$ z=>OQ*PKGgwHeF&G2x0`8d~bcY6m<@mR~h=RC*w{P&z^ z8!lsN^p9*@EI9055*&GRd3w{O#^ru`nUE#^%R?4?g^*=Et_)f5RYG>OJ7VMNkOf~8 zviwfP{bPNuHLjCkWUn{uYrs6$g)XsrzKQil<2>Kkx*=rgulc{L-c;D}btUebjq`QI z?zlqMYlgjit8qP&>@6Y7oJKV?`ad>q6CA$Y9^CZJqD^bY-!VUGoX?5u!-1*8hXS*Y z9t=!8j|8S(#s{X(A28g>?1S$&>^TSDXV~+@eD@m8b$^fHT=#by_H|&sy9|4d@c&N3 zx$f^U>~+tYI^rCD%s8(P@Z&<3+PM3cW?Aq(h9_nBiNfxaVK=WE@vCPYkZZ*;|5L&y zXa97MHv5+6gH!if@9*L^ju_|vwmehuMcosgoIfS53y6>TRQ{hX^7M>|i8^~WIQBK) z&E4Y{j2q-X#CYB?_2C~qC+w1s7X`H7~d$k#>j)e zy;*QMPj3|*=dbI%U2y){Q2)oqI~DHT3in=xdtY#@!v}(+E;XV(JebBEtl}-PCX&n37Klq}s!+GG`SHZCl{4;#}+PF^sL5yz=j|bTMH=&#R z0=nM?hkvdH_q}oC+du!Iu)|s6+>Zt4pJC^x3ioq``z1K#+poq^mp<2T6+6FIxIY9p zzzy2?GdSkIsrlokf2Vn_p`od5#%FB}mpzny`j?o4wf(!8!>th4KgMChKmS+Aa&OIs z7hBd+E5oHf$hNLtqGzGhlP_;BtM!O<`Oo7mhKN;as0nR_wMdxLk8n2gmxj zdFGtPIIb=H^Jxn^Ij(jUJJT7*dbpkG3p@T9dw7NlH)DmHslv@%;bsvWy>nKf^&K7%vErC-R3U1tT&I~hPi-^c?HKB&nGy1onLVD{6@iX z9xf1Eu4%aT#&vFXtLfhb4dv!im6`X%Y&2<(W z*MVIMZUa}0QC=sl0b%d8>!hw>ll|2*>G_$zPQvAx-z#yP`MHM3hoAZ7_YgQgw|gf$ z<+;68;yAbQ)hX=w8QUkxb}Qm?ir(mFRo^7rUC7d#{jAD4D9@^#uiBa8=Sn`~a$Nt@ z*|1jfwLBa0{5>jh_}P&2vz(BnA9`QSF_(Qc&sp}>9GB{y^oQkR{#-Bg9Pd{-SDU2x z=^a^*{qMaa>usF$$~~UxnZHZN^o71t@iP=|M!ue_OriyZ24UCJr7rFhD~^N``KSUocA6=A91{< zVSo8>>~%?PI$qnbzkE3M8lsLIuWQ&}J{u%ZFp{FZ3A4zGkeOWA7XE z4#yiB_LmRGeeB&|J{)gk@BZ@P*z1EDaopFizkE3Mx})|SdtFc)jyE&xFCUIKw|9T} zaJ+@R`^$&pE&aU-HDli^>=)aHUf0yNtt&WcZ9Tyak66|(xJH7s_g)2;^Rz+1QRlASyWsq@q5hAJ4J%xq z3b#>(+gNa{LtnvB^P5z-O)K1Hf+N?P3y!>PQQ@|%a9ah}x;cga^&p4a7)OtAymeuR z{qNki!LbkgGkn|5xK92-j2#S*2iW`eg3eT!W3n zhJQXp$Z{?WH7?i7o*|q28?wVH-0*_qT=e-z6r6vC&5^<78iX5VT&}?}hI0*$uGk$L z9BYv0+{-x5S^s?R!cO*mpNgG*jU(r7XTQRZe@1TiuW$!cxB~@8FFQzZoMi_Gm*a;! z)HwWd9}g+)XgN_X@*BIP!#&!0f zVjOQc*VqZ6+X?}DNB6`Ecaq@n_2l5PKlF}MjLS8Dn&IrvslqP0rxzUgcGu4+IRC7_ z|6}7!!IAT`3XUG{dS@3L*K*FCQ*izn*JkGmj%%p%3T^{ej8R@=tN~%~wQG#?!zO#7 zXVUZYa-D?B^RidsI4^T8kPkmE{d$u*9KrdS*V{jc!I`II{e0?!puJxp?9y|*&-YET z7YbQ=srUIDXW8d-PHO$ydvX~kb9i4}EBRCQ)jY>V6+7Nna}F*Rvh*PDWjU_0m*qKC zFJn($5_9CWWuJSGD8Eav=X}3rZexcr$~~8Bx!gN^PjyxH4)wuY&Ru5n(JLKa8uKH2 zuzwiCWsd~dJF+=OuEj1d;&3&-pIjk0u0yUA9N$U(SpYdVj^M5mdaT9Ofd- zdLjFUkY(S#S;%_6C)KC?*%g!G-V%D``fb6H+jlD5yMiNk?+K1vy)QU&^MT;V#fQPM zPdr-u{K&Z6!#=Laep1MODrCvcXCVvzT*%_X7lqv~E3#h|vR?~Xa`{cjg1;5AW z%RTImLN@m>p7+7p{8h}EYp2!oj#~cRA9{ZaJ>vLBaOnM8;Tm3=cs|y)mEg#6Yr(O0 z$K5`$P4-Qj&}-`=Hat(vH?48Lt|OeVF^%EQ_70xfa2LZ|$F()=YXF|gu&?!Sd-<5s zu&*_E3d21Nle5VU`{{1UlX$I&Bekl_)Q_p zUR^w7!S94D{m?T=Y_nB*(-LOK*Cp$buO$VC%?=fADZvr{(t;!2Wduij%Lsa`PsqO}GXfWY{_Q!NQ*rp*uuyV}cu6;r0|9HhuoGUb$veHi^H#h<{l4vO$GY zHb;iuFyYI8>ty?^^pp8|R_0SRLl0ZCqBmUlLj5uyXNP;%=>OR8^=-6wfB8`Bi5=w| z@s1GlvxiwvHUG%}#r)W5)$D51zvOgFksIB_na#WJ6%kVyj*Qu*JSsqLbjuE zoDbI;m%ry+Q^=a=f7is<8JBDK`jEvxmn9E32yTQ6*tk(}tjkS;!{#`_H8xpnXt-H$ z^eXj?DQv3eUBMkE`Yzn@f@A+5E;!;-HsO>_&ZMJ+O?>fZB9*<5^XeGEjfkAxB60?I ztKc}tm7ko)ienuVNA3Ibp~@J@r^*?eve_8%-6mo{?=&$Vex4^d{PgD|m4339)N_7V zqqBq_-0i|o*3X|GRnEs6-68Z?BNYSlDUQ7ObEV4pSX-sX8L#p}T`N86`Vui8dUuNW zh~adh$NH(YMNd7aiTUoT%y+lo8a-7u?hzcmsOJjdi+aWpoO-qnoO-qnF(^Gar8lBE zgob-X3~=`ej(Sn^!KwMki+a`#`{=!A+^Lq=R>Xe4@C82~5FC1HKJrrI##d~W^k z^*kGVQP0m|KRhVrqZd3RIQIL)f$2EqUkN6ZvO*|Ghx&HB$ zhI`z&e0}kRn4dg7DeS;KRoHo2*daHsMjy%kz&%sgc~;mV|IY=N?Z7=>*m*(Np~osa zaHX9Wg&o#T*?}wVyd>;!-n}UD1Lt~qy@&D`Z*w_X^qfg)FgsP{@8LWT}PH?ngqFJznPFVe%9{K)4%tzn!=N6TE?8`5O9%s}yLJylNKE8kY^Nh;*$eGfEyKUUWdiYY~ zz7jED^J~F1`ZPAa5gd6@&+NgbdVUaG>lfUzzK&=kIO^J;^;7B*yRr$VY*O?82%GrX zR`@~`QO2F0VZ(NoX8LGS07uWi`rXz%pMe~jy7ctmi28}{GJf&Vhx z#W1=0({R^>|1jLmF#nzQcfI`D6Xdl<&gUk&HKf%?U;_Z0QFXVbjg@^8rhZA@u6$2NuG9NXlEyO!`Zjl3}@eFH9V=BSi(53U)F5#z|_QIfvJf_1CMZ%HWo3Q zV_Vp8j%^{sIkp82cQrlcYi~Hmwt(RrTchC|+x&)eZ1WkORNgxqmwoGGIQ!PoaQ1Ck z!(Eea%NWkSEp0gawv^%QTL;70wZ8gK$x1NTxZ>t)1-&Ec=H7@VXO#+kmzJbZh#(~NEMuu~2eGKQ=HZ+`L>utDe zifseKIksMgb8PDy&atg$ILEfG;YsCvJL9r%+ZxWkZDTn5wzc7|$+xWxXWzCooPFEE zaQ1C;!`ZjZ3{PrrEcimpy|H~owtXSHv|+cK_dM$}`K1$coZ~yh99gfui`(cYIP_K) z^TDkmIBKWAu*vn*fC{HwCj5chLvUll?!XE+NO0Ix^U?d1P2wLQ;vW{i^saEq=E%?+Ec|4hc&1pc zQRZ8%GM}m$z9+3!(HkOsp?;Y!pM#ANgI?d3dlEa!H{u;C=I1(}^;Glk*`!z2m!2Ox zDyLhC+>p0n!Z*HO4i_9Tji_)V1xKxo5*+nAT5#-%F@nPv6(8#|R_L)Ol+BT`_9`#* zH8mgoX)iG!vF|N7Vo>WxjqD@zsL_1|*Dqq&PjJk)zu@@)(=7T#_aQtodVS*#R!v%+*w^sbT zt-`$_VqhJP5c9zuS>cWn9DY7h@#WD9*X5y!b4{;5TFeJ`jNs@?#|D>sCfsqx`S%g< z@xl(h?*zfI=S~!Jz?~FxfKL{(%yCM{f=>tvtNpD!|;udOdM+{JkM=>>+nCVal(ZicyTJ|7Rf1z5T=YWA`5eJrUD3Oy!i{-yVw=Qpt%e_N>ZxS)>?BX`YRk)i4M;$0VIHiX#msje5=RoCN$$q>=_yTvU;8?TU z1cxuT3l5ui2yRq!${Owz9M`RP2@YTG798uK^hU;f%1^F8?-4eM;ojidhMnB&A2F_z zeNL}`*l@1bhYaU>eb8`MlgHojhI73>U^v$+&#S_h`;E`_dY|F4Ua6l)gUjE+S+mEC z>)@&}o-pj^Km6kr-6sp(rwiSugf9L+Q|LYyy712mU3_}J(0ws<;a{lezEtSGQs}-c zbm@Vw7P_y8F8pgkmmc{>q5D?o!oMkW>7j2Iy6=WA{5yrNuOE4Q&p5A3WIr(Mb;-Q% zhc4HGUW4TEBjdaVvGrld(oekpko~x@e+t(^|u;V;2a1AhySb^NEo{afK0 zUh#=rekaFfE5Qv2KU)iq+H4aXxu=)A>g2|eKgW{^JM`x%D%_NU!|$moT-ypawcuEn zX#_XY1F$iz;LvMV;iju_(+iF>bOyl@`;5Vnd-o4+rY3Ge51F|lJ4+!uYsluk#yZYc z;byOJa|EY+owJFX;Okt4otr$c82-$$vWB^yox9L;N$PZ-3OBFdhos!eT7@F;E3DfT&UpuGv-@ZaQL~1;HdRQ1-GXgwy~Jt&|6$^oP$dQm+N3?v&ja3 zw#(A-v%htJ(f@7Pukd?GF$eb^?I1YxmJ%Fi!_tDo<}!lA=CT#8W5MNGo=U{WnszGm zJQK{D1zRrT{AKWH_!_Q3xhcBB6j+|{?;kFPQ^KB_O_O0U3+p5s( z8?pJDI@j`fkiT8H%xG> z%W%PQPL2>9dLt{`s0ue)aD4X}BRFCo8yva!9K-F^#7($f+PfmVPa(T+$hLBWF<8g_ zD%}1R?ttKwuLm}96MQ|Wu#?Z09c=%XJ!W|LcyOWTA5*7?RJcP0N4$p#4x5Jyj{F}{ z;f}0uM-?2gc`Qd4oPWlA#|RETj};vCe4OC;esjFw&^tkJoO34zm)Dx}YOt+e{(i$V z8gi}S_eo+7uGLQ#9D1h+j_)_83J#m62@ad5SGY3@F4yarO1++0==pjw-&qyzY{3!3 zIR(dg;h6LHcVUIQNO08F#e(CUxI}P#H@H-AkPfi zDteb!xGMz5S)i^Z*dM9}xdypX*klc_5*+Jrb%nb|aGdSe7Th>r+!+4spR(67-*tta zOA`C_f+L0-1c%KV3oh@Y(}XY7=uL%Q_GO&l`n!OQn+1n2w+N2;ZWSD{-&Wyn7aa55 zAvo$qap>Jy==E(5t?6HNUgSQU*S1&EEBi#Ysd@q58}2H6ag$sRo>95}Kd-{wEo}0) zV)qDc`-tOS!Li@(6WplKyI*j`_kiHgvU zJJh1G16SI4R@h*e)&Uf7`~UI;Gd2kynfj_19clNx!c zkbPOm;@>NU?5jc+A6_eDUl+2Y+$|e#6tZs$S=RooLiTMTOKk5HvhNC6YN53Io{(ja zmw9+!$g+1nD17)ZWOMJpUw<}B?$eIQy?Rayb9^LhqUX=ZDQ%AOk8FG_^jJTCUQVgU z8hFi=9P|0Jb4rf+YVrAVbV@z){fU^5zUj}@DfQTwp9(#$z5Mw)r5-j_eEcniKWnGd zBWFsFHTq2WLS8->9J&5NaOm;;np{id>`S3XUcM3>v8(xrPjS@l*TN>;H-aNC-&VNq z1UJglZ^NH8Q?3IX&z;G6A)nt1n{Yn}j{czZ;FKP@R(fzskA39N=qcla<9R(fKGx_* z5d+*$f@6);J=B@+XQ7A9Uj#SGUAOV8;K=oFg2R{J1xK!x9yPE0q|W{jHsSsh+^FUd z8vYU-F{qrO*XAWR(DL`=`1!Z6NlpACIC8D#V~wU2^TGXFvDxtI#D1dJs=~Dv96h9s z;OM)P32t=cdUC;06H^42uXpfsO5>I_8{nxbvTY05sf8^0o+f0$(+XMUY**Nwt|B{q zAv=SRC6_aXEO;g%OKxWlS@0}Ec68)=){q6y7P4)E_k0rf?8bF6jO?6-I~!)6IYO7U z>1G^xncKMT$<|yUOYh2mPdHCur;F(k_k6~6O?KxkWb@w>&SQI_N0ObtkTvQ*{~C?U z*W3$KWZM_A3kq3swooCvu#hE3ixjepR%90|WEU5*?9nZV(oh(a@tW(IYD`fF+y+U?a3fT>W zEOpVlklj$ovi5xn*^PuOv29$)_7$?!LTPssAY|^slUI|~w*|vfk5qjGR4!!Lw+zx_cU-Szu_XF}W zpoyDsZP{PQ67P;7i$6OFS-4$7mK^UKvUwk1ch?FxROFvr?I!eyOX1ALH>WsZYG z7JNv^<~_wa9xCj>9Ts-LhYLHze1xzAcVyTB9~H7~EoVOb_Xxywl5w33Qxhi|_TPJe zPcYoYFxPa)8}{E(fR8iW&G2x0`8d{acY6n4_e9J8AN?5P^Z(y_wBfP_i2vl^axKA~ zVq6DTjd7~5!+NOa`*59inwSIb^kR-PDs!AEIQINmVh*^oi#g5_cGz3z3Xbc?^TZr* z=NEHaAncH<3xyrl`l2STp#k5ov;FLAm)8f`a?Xn@b6z6mq?a-$*R7XUWXm}(6LV7c zmkT@O;L0X$!rzErQIWmMxLk)WTY8$veEIXd>m zjlvFl=%&KXI3Y_;ZZ2eR5wgT`Yax4^kfnxhFJ$i!vefaNh3s8Ime}qtWbYBO)M{z> zULi~Um3g>N$kIRVFMN1F$da@1h3qM@xAK}IOKoJ1xF3u;GRIzeNN`-=J}fxS0QKB4 z=2IN?{)n*2_3+z|HTRC}C!Bij5xzVsY~qW0J`*{+TwGgWvo@c4?h*4Tj@Ta)eiGk% zmDs1db>h4*pL)Iz@!eaAPuaxJ$AzERd_r)I(N~@n9Ch|^rFPXbe3@gJWs|+@&)z9*;-^37rR3;mJo6>j^#~WY@s!9J+|z=i?JC=)g~Z227bx{PbtOlpNQ!>e(f%(OW_f?pfg{>*vqJ zDd!vEAK7?L=&?pB2If;7dGTlEl=HE+N{@3$<%POddepT)7pI&Lz2`-I#K7~3a$U22 zYHiU|&z>Ql>Uk*SMREA~g7A~x`t@TI#|NjL$H4CsPQQI(J$zBmjUc|cUzk{r_|$xy zw=aqqsCj?JL%Bw9F9|(-d0BAiy;9*`6&%;euL+KIcwKO#-E|vp2#$Qd8C?F}fS+#} zm;cV;?TYL>h3vaRmfXA-vf%fHEIxct*!{2~`%xkLv5+N~pM)&c$dc=MLKZx4 z$f_PTpK-Z|H5$%6Z2r(?ZE_E5Z(R3OZwrJhwU~R@f`y&j!xlC!_ppTu+1$gHDtg!= z#hkfzI#>41qC$^aSWIx#!r~Qf3Bj?pOA3y)>>xPSZu+}DqLzP;#J*W7^xC?J4bKzv zbu-S_mFuUjhJ9VZT@3pgaDCj_u&)8QlVM+Le)rtbu&*_ES;IXHle1+E`+DQs(uUpd zcK&9=eM5G2;~X>BYJt)18JM|N4UF%r7|yY+Y&gfZlHnZNiiUG+D;UnPEpIr-ww&P} z$=@D^b8OuWdu%H2y^PDgt#3H{ww~ea+q#CcZ|fM&zO8LI`?i+h9?7>g4QJohFzmju zUmZ32KQprQnEfE5Wgcw+=4XA%1RST&}HcE3(@avfB$;aoy=pWhGwW;4tc+SpY6}zil*z!IghleaR z==NB*fufIcop6NEBY#Q{?nt3WernI~>TCJ%_$TwX82#h@*e073_TVvr`JLHOf${xl z;ltRlbzE>$GKsx6#?Qwa*EZo3E^odk7QAD>Cd=4=&Uk6V4A1>|1C8;)@upv%P58BO z#NR4p{f&4|5b-pI+=+r47Th4QzO21^rWWhO`qCfWT)7tio805i7~#*T@bx4S2fuq+ zL*xcu)w7O>@nm5WZg1fW`97thr(&S?6~}z13Y+aC_R}id>4GDlX9$kxdYmaZ^v)6- zdS?p`KhF`|fQaE-!C~{f3U|KX8pGxVg2R^!1;>0p-QfmWUIW7YT;V!BGO-@#i^?^e z$|tpZk?<4l;^10Gy#5AviE)0N2fnmo=Q6>K44ao%xGRFo^P_jAaXvrzDq#oTt`>IS zt_eHfYlR)w=DG@Zy_f^N8;t7^@%r11!VdG_RN=;nIncY=xV#>>2wARiZZ$6Zv+m>W zRm-y-+1m=)UKQEf3)zh;vfDk~(uYljEcJUwVRs85OOLy=klk9ya_w|iA-kQBrH^0v zY)gN457~SkvfdZp*_=foA8Fl zZ<26t!>pr!yn$hS0+VOQy^QCqb-cb|&r|k|8v3yCZA+8QzB$gmZDBb3wz=W#+h&Hd zZ;rEXn;Oo(ZDP1wLvlJ^!ZR?*6(Sd z$9(F!Pv|`pdik7&d)Byo);uTXAW!P~E%c-3#T;-i6mz^-nd2qFvA!>hIpAI?=6F@u zp}t-d+!!wp8?OhKa|QQ?aXD9ShHUP$ps~T+Z#MA)9l{{`pMUVQoKe;wJnq?2C%* zm&WBezY5trC;ReiVTV2RO=0I-AxloaD`dYHvc&R3A^W3{rT%{^WPcX2W1366;g>@8 zS0PJmzZJ5-3t8&AwEKsUrT)r1{3&GVAAc1-{4He3**}G>Kew#x`^f$qcJg`?cf;!w z>#>(w365*a))lTzaJjdjH<@v{7ALRRnWDl?S>dJ<9PzddF8c;IwQ>I41w4(gLmbmq zxOQR=^rkZ|$3MNW!|#k{5FGn!#tJu6g_~J$toJN}<2pw9g5In`4{o*!H+zNK`jv_6 zl3qNA(4*GpEI9Jz*5)cWVsmcpg7eR?IgjA_xqyv%1&7`FD%|`Pu2FDYdoLh3>a~4^ zTd=|{RN)pD9C=wqtP#Cz5updSsNe=Rr?X)(!J)Uf;Luw_aIF23f+L0wf}_snt;DyK z(8HId1-EC!w~XM>TeiY=6ddt&5*+b$798_+5gfVhTH(4?xbA|ZpY;$N^DQSh{9Ime z{Ue_%2#&LNMZw|AN`gaA^%s0uS?IA2s|XIgRRu?#spkx0v!~F*=4usA?P+Ro^@`pa zg2R_J1&1$d2@bur1;_p3*AW~x*A*Ok>j{oCXMMr(`}bagLvI7Yq1U^@Z74Wws&kXG zqmR%dXB$T1>SdJRusgcK zjj3>B1xGK~OK|Mly(`>46>i@Ox1Zq1i+Tw!$4(;f@y^J?jL) zG2e-T!_Sih$2IiHg5&%D&#G{0 z-%@91SM<)QaOVmRKhF~!*SqHnju5FFQ-s&*UwLdC|NLT_j@)o_>K zuz9!OsIz+n*BCbM6&$szp1p%j^&B8Lo)^<9=I}Q-e_l*^olH$UEI8s*HsO>__NG6- zr<@O8{8=$2H_AV-3vNua$p+;o|E))Htb^j{XX<%7 zKOYb=pr@YmfuH^ynKGaF>CcTRIeL$J{t0X3&(bON;KmC-SwDZSPN~NlJt*{8BNYSl zDUQ7O^LEPlSX-q>zgBsnu9Y5j?a$UJ=R@xy5g#$|jH6uFte;w2^wjffnD61re2)mO z(OtLksNnELJ#PhH)U#jU)H6rm)H6qjLFvIMJ!v4d~j+$@}iyr!aizv z$+Oz>?`(NVJ}Ef#)O_Tn#ywTB`IGov3Gu1teBg_Et_b_#X)zzY;2FWO-=7s6 zdQ-eQaX#6%&j~$ZcwTUPe^B!gpW>*A7s96c@5wJ3*TG$l@sgOIJiRRJz`at~c~#gU zH?u}=vOjRI6?R@1cF6x5!DTyeZx(jm5_agZ$_`v<=WSt!wNrNBN;~feJDhhjiu}O2 zUS6+vg&k_*z2I_w;NCCnc;0>gXIbj!gF^O0A&Y+>6|x@-S$z1Uko{E1QWu{UvY!iC z*8Yn^_Ddm4Y+n_!Ukh1kp|tyrkY$gTdH7byvUk2KeE2?O^Er#Z%Zk22?i&Ti96tz~ z=(QJm&a9V^!gt5)=S3XZy7 zPH@DoY{Ds<)clIVCV5#~aMa++g5!GUcabx=KLp2qP=3-^6vsL!j(xO-@RNM1oWUuZ zjZGEN(D0{-0ljs^eE8W%aQL~N;JEHo&y!${dI>$azl5KxUvHtu8vQNwSR)k!^C^zJ zY%JzuZIvE9OXY>SR(jO+W@0|{{t@vJ!@5F`^;2t$o_Zcg`>2*<`oeE7Za%Y3sIdj^-U?^ugzZBGB2QgFNQVKO%tgZi4jiJS0uxA!4p z@9^K4<}Btoz0k{c1_?WjAwO3!zkkeHvwpcQ`iI`0O*ZXWB7ye^Ouhz+IoUJnSsm=P z+A|sUh)^?Yv4pV0{7Y824q^`UmNG7{$I?QUeYK2n z*`MR>`3>Hua-NZ0wvaukBHOW$J)?J~$Jv8X~mj2!lvU$I;htI3*Z-0J4c_uJNSM#Ujh^t$LQ#Lufy9+(I9)cr=*ZdM30A3as|pUAYCf(Zl}+NW z-JAb9ksHosWs@G+Q~1d`@qCV4huCa>u13r2Jvdb}?D@+oey%2bp?;Y!_mRe?)t%J( zw)71<$~WR&UCd9PU_I6RYc%PV^`+;>j>@S&N208K^0ub%jsCor;D~AM3b&5nsI_$k zNByoRIQGQ)g2NXTAM4Uf=&>i1P3l+Wg}$cdqd#pR<|Fprf+GgCe$>c@LXR5l6I`wz zV%W&Ie7&);kR_&zhc;HM%nFK*xe-T=GtRTAAYIXPUct}#WDY;!Y0?TcZ+ij zPTljpzq?@LtlRxxy?@lb=~tEOzAncGE zzi)Y2`*1rJc6JhW$p6m4<+$K>DeUYj?9iu_9k|lYZo&?0r|iI$c6JwbI5Yfy@nwGC zTraQJ9>NYaF)+BCAGkq<9nX6?Cp9v-kR2jq@o#7$yQh%Fhhc^6a3M=wj3{JB3R%{E zR3STB$P(L_LUyc>r4~xNdkI`MEa(1ZDBQJ*uj@Z?F#HToF z_i$ko?g+t=mm@3OQGz4(c|`1R^9qi99xZIb9V0mUgVKXjdgNN^!6`lV(HtTMxH$#K z8XYTQfICibtP%H2RlVnUp@+>A1jjm@C^&L`lHl;=WWkYZrAN&xKdG}*giW|p1xF00 z362<4&d^&*tPy^mE^Ja0X9$j5tNEyjZel*TGb=XF5*&JGSGaQoM-MqyaP-~t1V?W@ zUvSjK1;OR(9sIn|IRCB&zNjL5aUpw&kR{)jhAj9pAe4UUb*Vl(E_=b?>cXggm;=a+iPKJ>kXSlOr=D8_!S(|QYEp9Qcd$M(N z$kJz)OR~2XcDfix+_xLoHObyq$aXaDzu&#sUg(iz?r3)%aGEVXlgA^Si@c6=fGppYeJ4;8Wx3t86ikwW%SA(dsWZAb*6tbT0 z@;VRMCk023eoAok+Py}wq2By6&JxA3CjM-O(ofd*H^JfO(-l9T5gdL# zD>&lYTExd$@|@7)EO}mV)P|am_!P$)y%09p>z*aJ7mdq%=_N5gd3sscfqSK}^Qy2z zZg!8{WPjjZE9|^3?2!LAg3ET`-Yo3ACG1dB$_`v<=WSt!wNrNBN;~feJJj*cB0q4h zm)GlEVTYP{FSwi^xc3V?o_F6nS(Y04ppgAg$l~8eh3v;d79Tz-WIq+M)Wv6o?B_z3 zwf~}!{Zhyh+gF9`*Fu(BDD8eDWZC0o9=;W_?49olAHEOS-2d?R(5Q{fk^7N?V~!t$ zP4q?!J@Wpe&}03^3cWF2MmD@=%3jQTdka10tHrmk&?DbJiTUW8`wKnx<qEw`5jWS4{&{1NiClyH_W_={koySh{+rF2ulauuAJpGW{n^A#_?xLeLYKAqD=_Q$ci77F zQXBt-Ecjm`i;oR&w)C-;;j)e!o72(II=DRNzQH|c?@Pze{`UM8{%_0w7DGPCSze3& zE^cG;ra313ol={?F*i1O+R$*k)qRfn@Yn}l z(+qr@BI2O8P7|1UrVPw}n<_B>kMz`m$wS+~d&NZ48unU4zg=MX>6+wQ#)$2+#Qq&0 zcBdD5?Sq>^aOlk_xB;Ozli^ zH#X-oE?>{gV>tJSxrN<<;p@DDBRBI^xcMtwV})BlaQM=`!YvqFjvqaL9z%|QA(Qp^ z;TINm>3NG(xJ4`6VuJhc@8AT-S}jrGmaK3cg3Iw^(`zcnzqH}J{!0nF^ay`uLfKQf z-dRTIjfuKh`N@g(I2)D~dgQyK;OH;PCTCTrLeFcII_+F=?iqXk%kwAp6PsOx9)9wy zgM21n2y!#NjwO!kj=ZnEF< zI)J%)irA^8)hgWTf+L^Y&pq!==JVLfd~$uH<|CgCuTQK;?W+0Uluh=-8p0Q1;C}x( z25dIpr@v)B;Z%Kc4yb$dqqnB;g}uUjQ+f(wG{(ECrT?#M2ftLTs7e*1;<|KB{=ri27+Th_ZA$!sQ6gt z4TTG=w&UDK_cmPC z0Qpn3L%n=8;e-s`c}A21jqI5WR+gIsnBD-%_`jHf@8idD%_UAwe}RmKyNGK z@^$gn!VbP|Q{lF)aN7xv`KJ?k;+j?6PmAxe+Y6g;J5;!Sg2R^?gfIOfU;Tw1+<*$V zqu|)nvxxcFV><~w^0IS<+eL8L+*NSIzMJ5PZ}$qfM}-?GIM!%Th3hEt!ag1>^x%e6 zxS@h0zCA14Fu`%I3>O?WM+gqRO+|cgBZVG4d6eLoPsNVS(L#^&VT|C=8!I?^t=cQp z^U=(@Kg0c{5-0nr+Nr_M+-f2eT?AnMXft}#|k~J-Hxkp#|w_0 ze1hPxxrW&9a3>1Av2o_8J;pvdN$A0yEI8`y6v5GxPZb>072C-k_MI$v=7KjIe%j=H{3aKxvc z!9ef1Na(Rf7YmMhxg@xyBBu87?r+Z-w{XHw8}6C#Q-)Vd_({X7C;WurH4=W@@R|uf zW_YcHA2qyo!jBkUC*g+;ub=QkhI=Lapy3S?9&fmJ!VehUFyZ?R_euCZ!y6}jui?H4 z-(z@_gzq-IS;AKs-ZJ643~!b2orYO2|M(8WtdZl}4O4fHZ!=7dI=*n|t!+b|{ ze6wMCgyV6B`7Y@ACd1tK+3}5r>BEk1Fied)zTPnV&hd4IsZ+<-8XoFD#JI*V_3s~D zZJ7S)_$tHHvg0cabM5Nb^G9!Ve7WKA{zHss3_k;~_sa}FoA9NEKS=ml!~UIc%Jf?*kpTEBM9&XJ228xmzU}*7l_e_i}}MMR3%I>QmI(t3nU% zwF>vT;HZZ|;>=?$-w=A#!J8HCEx}>)ZNU-yJAxy=cPrd`74Ch(v3?&^xT8c~*fSpr zJ-Cl5+{c0=zE3LLr-Ea@e%e-a$?{VX{A{6%nFTm33HZ2l%VeED5))UK*Q zeEFlISF6uID|&wkj+*#eaQLZu9h~X~#PE->Nk0D-95&V3qSx@29}F%3`vzxYs|u%{ z9f4kJp~rcuYM#B>M(Dv!COB$ha>3Cfrw|;mtDX!urO@jiXXsRdL$9si$jj7%LvI?v zp;zlY(+WM#qjrMhJHvE>KPa8_nCwqxt>{YTszGoIBd=; zIDAphC7|yrj%(N1giY#tsi!91M{w%C==d_bu!%3~{_WKK$FEF0A2w_Asr#ZcpW=vp z4&f*9EnA8G$-5@@llj#Bzlra!N_@&De$FZU#O7RrYjpQ)%q=+fgSxLaHq|}8;kXxc ztL7=1{)O{zegbYm45puTDH4`Be8yCNGM^&-TL4u|AEBl`6ee-LI8%`m+06yXAEq zzNq`k65p1U_?T~UH*LeuXwK<{EC!#O`5pVU2G<*l820bpeE(k1u;-23Eo`{-7uiJx zhu&g>Bd?1KjyRVnxQ;%3j3o=M6NJ5&oKx8BQ0QfDse4e=hhEl2aLnIT%mLRe<^XpWvdqyVWWmdYY~HtU%L|UU zRuCL|D+-RivQlt4S8yvE=eYu}BJ42#s=^Lj&#(htP1wQL)q`srYv*-K9@jU{Yn8pR zp5e|3uWPuAVfyhphPx)bwqdVJ_T5^ByC=NcGtK``GxfKo@yi)c{A(C4_b`6c-Z;mhU~Zj0b@{OB!d z|KCrJe@m0idv_~gm%g%fh1;gWZ7VpwZ*C_z)?)h#w?l>N7hH}Xn_g2n{sD&b`u7)h zsa3zXYFVpXQ|u`8xGviBsfqQthS^Ezk?)-a$G%lI*<-8N8kc^uzjqNfIs23z+^#~8 zywvU^J1(vbCUci!aQ+W$vN_?u^Xw6r>w(<@)5uNU^J*x9>cXP?5(zCxCsv|q@A_Yc`z@A!9sumg8s*a06T?2x;I z1;=_EA~@FLP{B1u9uE^7HN3T`8*HlkWWw!S(c49E#J;`Yh)>ysQ#Lts`U{)*IY@Bq z(H#ZHo;h6P4DJZQamFb>xn@)x>!3K!>fME(E%aC;6$A4rj=YQ!^Rc!{kG`SuLR~98>iYj-?kvOP zDuTAXNwDDVvIK`V5d}k&Mx94KXo5GVo~?q<9Ik#jK^MZnqVC7rwc|u?n{@?=^VFb zh<@a7reM4$QR9)HVyua?!l%0aIonuQpZaL$i1De@xxx?Zyu#1N%zqbloj`6Y%w+UO0o!g5T?g-nwcMz}NH?Q2Msr{XTF~*(3C;j~Xg{4pG zewXOS{Q12M%YMv(ubGlDp5NE7WQ;D?z3&#|v2Q*k`f*&|Bl>Y3dPMZYr^;`X zr)ACWgIMNA%~U_wSK{~b#8PATiX8BHpJ2@%p0)c0qb}-RfB01Q9fWazy1aH_et){M zF05<6M_|dwUHOD5pRD;;gim7ed()NUu?D$MUCx2)od-nCU}FX2cu;Y&uPDYmD8_N5 z?(auERn1_^XS0{X+Jhno`f*>eTo>Z>`_`59B+d^7ivq91dla$r2gsEgmOxU8I4cnX-YvVcL2ljmULBAmUP=gnRAJ|Lb2mNx`c8oW9wDxLh`o3U` z*z3)~dd%bhX#BikOt2q{e*QA!{U{hU{;6Pjye;~;m&WmS_Lr>vT=etT81EOssOPU0 z_FIMhUSWR(%j<;k{xs&t2l_8z%iPaw|L=J|mSz$EFFLk=i+K8b5v=`FVNGwlMPp6r z*P+5X3PvoQDr~%9xi;*1evdoP7d=npn{63mp`te9i!qu12`X&D3Y$nU_M*;$F{cw( z*d!G;X|S9xK7B3bbu*c;r8YAZwOLrK8}_ZqMLguF?o~%k)jjo^eJNR+qVVbKi8VE4 z!Q7tnn7SX`Qf^>vD$$P|rWTAGYWH1J_nafYl|;>ma~d%ov8aBqX+=Nwz1ltCzOm!N zkI{S%9pEono3Yhr`~M#>Luh_yJzZ#GpI*ek@8xHzu&*PA4sFA>{>10ZqF-~^ttRTo z+^j4Z^;h>QqgK>d{m!0x{BQCno_-=8-jB>8a%7)s5w#&!b&n=;oK^UQttVok-m_Ku zsT^2;iZR~o6`yldSaXHVSz*UK>+pIh+v%4 z78Q&f788v3+4DTtR?m(Rn?DVRe=pwJul;Yty9qy>x4H|)`=>9Tb%%{>nf>QC!HBzu z@X6h&*IV>sK9>-T+1;ghGf)S^>H_^ZtPu+`!<9|irlk<$a*9^zi zN}?Z3%`HAx7X6rqK7!#>%@J%B(U0@!sui}HV0%T))Ze}`p1SW1Yre1WNq(zWn5rlJ z))4)u=bD0XPG3tf*3#O7v6lJ?#@by+FzRxsI953R))oCYAFd}Dxvwu6wcbE5#@kRZ z;@qghHWrM1cavZpV~zL|b!j`M>-#>}#A+{MJ_o(xUfcKmBV9(@w219)fIV+kFly^2 zn-|Pq;~I2}V7w0?AGg`Y*!KR9Xj_}+xbl~_GR?Wg^_Jl)Uqf>~+twK8a-={9%(wy&n~8OjdwPk=Xxj8dEM-2daOH& zwx4M~23QL@*8`2^SPlr=NuqDIJ;qqB%h9HD9gZ@c>u{v$oc9r?^H_(Q<~a4?4l|we zKGZaOuk%Ap=e!R#o%24(wC6d2KSd*+qt0oZ^ApWB$9970JofRXbIp%4o!9cQrhTjl z8OWZg#TmwO-A*@69BzJ^={)wSrt{dRn9g~fY?{5!`=4Yw&!;Qa=UK*b&CU#4j>lZT zvxBMQ{(NIOhx1J5xjWZ%uFpB)YkUtT8f#^YG5+o7V$*p)y~uR-eWB@`{{^OVzWLa_ z#8}=3FAdvH-o;w>dxf!Fo6AjeeskN)qA#^t$zP0iWw4H+GrP)I9}Ml8^8>9IOH|eWQdrgTk+xniLW4_k(YN@a9 z`5Po%?)jT1R_^)PkMH?wCtKh1vtRc7DUuEI>3e>T+4t}q^UbDLPW^ljUo~mp!*eXY zhxbkVCey1YeWU3-Prire*l#esPO`t=bguVxrgLt-hv&KTJv`65@8NmP=iF{Fmd7sl z>}-3t+2%UjWjg15r|F#c9j0@hzGvsLZa3X8<>`BN&hs|Y8z%01cFxoH?40MVrajMb z?b&&beb3IZ`JSD}_B}h-*!S$b7XPO`JJ-$k>|B@oOy{wE&(34tYdYt3kLk@)9sZ|1 zJJ*aoyJPgvaoulzb3MyFJlD(j@Z8t;@I03fn9lY3pZ4&)j(iW#d!Fy%+4qB{b3MkI z&g(hH`H->oQmx88Jo|mbY;&z1HoZ~mTkheGb`Jh%#Lyx1_3`Xaj}`2Y#2zo0Bd+hB z2$t9H*y#KJ<9>!?Vyq{N7?N`T!lR^7}we_2*$j=C>S}sB-jv7(%Nay`Oo%iZC3;Qhqaem`?c5N72&6U zn`zUlg6$e?K#Y;+b40M61S9U(gwIjY?{&fG_l989N!>q=es79?%;#Hzk^A-{2gXzP zDWu=q!YBRS5sYiUcLk%Lx+fj^wTOPO_XHadwSHeP*3Som;qybmh*RB*aA1r#q4<9b zIL1E`K6y`8`^@fR(U0Rw%^yBL5&f9IPX)v0XMzon{64R+F9gGS}@*^d?Oh1{;gol`*(t|&b}9n+||7#IUdyK)a3Vr@X0a#qhR>_Nibshxx#)C zjJ@vHU>#!}_!GIe9b@&qcF*X?@jpfA#q62)`V+p>)U>#1avVQ2@3+Sb%^Y_MO>Pr} zX7BD8ntAVF+CRHOH<|99^gotk57UhIx9Ofq|7E(DY5f0bx_8olm|ns(@&9gmNz=sp zo9S^~8%$)el;<+$v$HX;6MAA{%V!?TJlJN7Ue@%u{7zE%oizM*@GjA)mAW?_c`3&D zlL?=!g{7bIK^m`>VCvpjthXyF_ncDq@M1oGc>TZQksssb^-aAeFX}W!k(2vpj42D| zdk6c+R0Z?%0_TRgi}Qw$&G=In{R}lNU#Krh@r#NFyHBWlw-Tm&}h7TFLx%&SNtd z<7GBWu)G%NH>)up13g>C&+LLRr*l+TbFe%<{pK{5_ku3M53$W9{J`c8Kj?XcpFO-N z*5(zAc;*v~bH@CFZP(_qX@Lq`Fj&r?ehV4v8b17~tMEfk3s=}8Vhs8%YAnz3V!{vW zcIh})a;zNhi(9`O>l(4&Dcf#^?b;RF?uG3J728oSHO8>9u;u*Oqwu?#uw`HGS=eqV zYbqUT!k%PVJlSFih>c#N)@(puv{CCnT6hHJZ3hp92eUb zwdo_qFO1>Mun{zEY}8~z83SkSxeYb|Fw$R3=r#vV{mN| z4>{KE^~F8O)cJjBn@4Y5+Mb%04b3&h2oVq06x?$v*MPkGMGesFMC{|mIQ~Rl>$S4> zYom3eFMR#b@NUNQEzrt`7c zKQw$mt9@g>T#F2>uwz9F{E)|p z3fobPLBE}h<@|RJ+bLpBe0=J^zp?z8z`mh*kG@Z6;@vwm>tVmptcSfqvmW*|-Oa{B zk2ak@6BuQ>hjGT+!*tK2cQ>6s6WGmk{!Cz3)A=)jk*58<7V++4I`0*(h-Hh)@uV30 z;{n1a$Nlop{&zf>y5AA894LH}gSsyh$M60k2YlAXJ4EzjJjKZUAQ30|tyIZf-DijK z)O~%(?=%r7`6-{od9a8RpN9z6?9*oLP{BA4s{7^OQ{B4<#ywutxxnx7Qm#+xa<*XP zr+mVcPd=0Pd+wCu5vSkZrDVKEss|=UHkoT%JJxTw8)Pft`hy2KQ*`Xd!n*0 zsr#T%7sZJ47!fCX>pGRaRo#by&xm%8&!_WoODyVs5#;wnB|kMDYxh`@1MIk9`ML;p zyfMEnL7!0ZbE06ZpOY%=_~AP3%nCb8j6uJ% zjdk_q9_<|Ahw;y?u=B(i^gG{Jo{urYmc8r(V>zBpVy)&nV|!s?yLrX-qQZ9DitVm1 zH^wkf*zOedy14K=SlDt*UQ*Z&6SjN~cWGfeT-b8#o%C8`yraW5AE(U4W#-rSwp>rx z<;L%un2L+sZJKd>9Z5Bf&ohrRzM z!Km}ig0aWlA{aT|Dj4U5+XUl0d%IwqOYabj{pL==@OhVD#Bz7AiDDl8iTv(qW$o{+ z#)hUg_l73t`$F?x?t#!;C)^);uNdn=)4ulT|4=J$d@oO&4-3XPk5t&Bf)RfY@&AGm z>uln8nAGL5iqFR@?1>6{vcjIKu%|2RnPB<2C--NKdF|0dp7u#;tUdGcTxE>sE9`{| zd$Gb^5{y`0uCP~v<$UQksr_z0ueDc&AI^`jRoG?E{WlRzfoatR@hs? za=!TVHIwJ}ZQ+NtPTbi~7f;jL9XGfA{Q>8wgU{z6ZuQ@upCjVr z{tvlMoTpl%-stxt-|>=ft_eP9W$mAzzaM?!ABJY{`zUCqoko&PWs{3xRK5F-KSWfK4^iw`L<_>+>18$rz*kKiR-jo0BM-E?z z++kl<*jIuP%c+%EPOGq~pZaf{WBes+yNH^NX(jgmpNRjL3HG&!lePAjT?d!-WNm#T z`oX>xjBBs&1nVDhs=weo-`y=O))sa8wKB$Uf^l4_ z`v|Zv|6b|$hhWt3Pr@PiCw#)j7mT$$fndZkp(yE zeUf1O-kr7LWt-d>pNYAi%rxVchdd^v?+j5(!X9LrNx*whs^O@&P>816JP>ojA*I3Ie~nF{8waV*VTVY5`&tQ9t!V9dkpg0a@-sIcY=n^Q1q-9<3! zHdlqsU19SCV-IIf_PWkzj6KEmyoDc*M`!Z~<2dlwh;0F5-TZ}U3z`0RAG6mMEd06| zeY*xDzRVUjM!o&@MG8NBX5ehmg86HVu~>yIUSZvWso1(3V_o`KJt}^BR#-2=*pqt) zWBeO^plJKso_rI+p7*Wn(MyOinA;_bG2A72EoF>j-Cti?*zVxK+A_xST3R-2JG9}} zpV%%}VapdRuh$g{=1%dsVz9hEVJjKS>$8vPygpa1_+6!7dCjdFjQPp&tY(bgP5JA6 z3qLup)hmA1FvfhipEV0V{u*^(tHRc7Y%V2rk!nQKj-2;!dwduTW zx2gEuw!-=gMyvx0#&w5>Ffdq-mAz^^WBC{xY&yp}sPNnBFY+2vVM7bXJbBk)1@qSi z_z!E_3&vdSP%zFD-fwupcrWX0M8W(u*4&PQajmyg!PfSU(aLL~J^*{JT?_3TJ~_Tt zO#S@4-Z$y;yuMaqoYy;~*Qke|*Zn;RW4MC(xw)Ue5Uo5nZw|0$&dth?pGVhCw!4Ti z+5h}}v_Y~RDQwyE{d|;bP@a!+y=vzbKX>H!VL7k=Y5!j}#aizFdHh`~dHDXH>$97% zWk2+NHRoLJt9i_FU(I=`y_5Z6_ZUB~m%Lx)T5Xu}XYa`U^4^jAt)Kdp$9O(=%42n- z)K49&~njIez1K6<2rI*!TS3!*3{q9j%g)L zO&^H-cwesKq~9IlGc96Kf0s=^_4n$;xu3{^ew{0Gw7=*#Hu5__Fy`Sv!N}nt!8jKm zC2}XugGE2sA;I!K+*s@O&-3`3U)HOqL!Az-jB%J?9CzyP?5We?m3~JEMh%Y?jQ#K^ z!FX?fbcLz>7*Exn9F7q_xz0USFnk^-7;~oX=fHk-yy!=*)%`IT?*!40b$6m*_&iB4 z_P3J-*g(f914?^&WRd7e}BJumvg z&lP=%>HMPa1<@BirqcJqqVL5;--|?F*7GGr-^-#e{8G`EwRw5b_sZxCze4oocfVH^ zeXohW@T-fyK7Z8lT4UWzV|%@6f4{@H*F|6Ut6pjDZ#33B`MM!&*{_yNwl@`idKjZF zw;1c0Y;P`X7dOVSahtIvlI^X9t)YKEue#k>Ub}aME%CW6b+}V7_WHX72H z-($>cjqSY^KlfGq++X;4K-hA9Fg9$_4~A{dnRp%&eqaxWAM_)IAOC-!`2U|CEtt<0 z^YvIUhM&Lu;q7r_eoS&)K2i8_n`ZxE?McCwYC}y=70lNMjlX4IofqtxIe)sfU;DM@ z^P$Q4nb7Q^&xK}PJZrky{Q1M%3#NU(_i9y5P4AiZwaa+#hUVJ#ozRT;c4&@^me7p< zR%ph5)3hJE=r>IJv5S7)bPv;v_nK)x&d{%#_Tw4{_y5{1AF#|H`g25vp>AK-pHQ);m!5N_UsRDt~aq~e|U4fsXhC{o9oT&*&p6q z`+G>@a$8=N~_ z`?2j0Z?3nuXMcEe?VlC!d4cP^cewV?I{e|ywXZ$apKCu>I9^=uY}y~*T<>Df{_y5{ zq&@q?n`>X|tasOb?D)f*Yu`ux;m!3P_UsRDu1DFkKfJjfZO{Jj=6X+i_J=puzBjNx zxc2vem4@!R5E{!`TDrq7w!2A(~H?N{Bz;g*Mzq(1!L~Ms<5vs?3)VvRxo1uuEM?#mh-3I z!uJ0n3OVrt=&8ctJT(^lAY!E(O%^fi;`cY5K6 zbEJx!y=(?C25iO{13i-%gX4Q^+4_ovcVaqwQ$r|-(4yI~1 zr?FhSE@7Mfkn>!^4{Yx6gPtdBJ2$np#qEi_7B<$+G;`9`boZneGTp;8*KG@$?wRxg zrhA#@Tr$6DzeYe`^HSscfccCsX`Fh@Yq~tf*$WpDj2bUmVT%dIc#8+i_2-(idn;?d zChI0_IWP7ITjJ>{Y+=2_mYVkp+dQxMT|%(trna`EVDwu`FwPrGSJ*NYwya?6IfE+K z1Y!5Dw1sFS+(=xhCQ zY-^bI13W?rxfK`}p03#AN3^un6TPG9(w2IC^l+Ok z`eW0jEq-=tY4p3Z=`z0E{3UDZ9vp76uI zeZFAqqhku@=SSKFE&i|ewb2b=&n0t&&kI{zwfA!cyQsxgJM(J*n)1n7ytu`!+WWb} z=Orz!8X4>KQo%SDFRQT21>-opB3NES)cUGc)_x6irLZOLtHYLft`WAd>%x}2t_|A` zP3;Z8*9Xh-HHX;^t*rfTk#7v!?1wRK5{z@~%@uZwVASqb!3H+9wPT+4=f-(#ZbMDC zwfH|YyS>8hsIWUL?5+yCTQKswM=;Jc>OKcB?sJgm0p|BPD07FYdma$Wy&@LaeZlfM z3g-8C=;}j8yT7GvUxYtkx{RmUU$Qo~#eX)=6R~buncoLRzh-x3?V%R`+32&`U$XXa zi~nqYKjNy{e^~SR&3Qa(+V^+bBQ5@~cE27cv3`GEBXiZoe^`60#ecT{ zyV}Fs8!e5W<9+sW@)^I6^vbNdl{pw(4*q;^lSB3p87&ZGx zFpkZp_nNG=pSL(J{nd`fsGsW&VcWq)H0HWfg^ec|`Ho*<6I9rQg7uG>CK8Ojud`tE zo4CRzsjx`}V?UitFmj(fSVxL!&#)<4S^MWhQ&w!JDr~0?+YZSu^Egd~OvX0Hn^`dOou$HN6^!v_6O6f;y~5_Gu;zl1 zx92%$!TdGG>mnF&&LtRYeQv>aamUu?5sZHG3dT8WzF>JBux~GAbGt~J`_`W$#at8j z{9=q9qVEEN(QiS)xHelzFno3u44(^E*dhhXYq^iek2zhm=;x_1-eQ6g=i-9VubW`> z>s~O7vS@&xd{c@aZ3&yo|Kfwk> ze(MNEoa+ikEb9r@KYXrVVH*g>cpD1FI#Y~(8x{RFh&cR7oxAcmA)oKpO8v@xDf?7= z&*jSf5f?Yg>4}i=ie;_<9bs4ttsne zE75OE#Im(u%0|e_IaSjx0Ot9?)bsGtXWRFKH`a_}Y^dpKfStUL90Xq;nr5o5qc7GvyM8DlrW_K5g)7h}No zD8?8i{4j5$1!EuEQ;Y%Is~BT%;fGr7Bm8iT?%T@R{~p8t-Y>u3cZF>^=6;ni_ZMTb z<`|Rx?0|}GIp%?4Opcv{gdb{fNGoeU4;);vJ=9oUhlho22QNl6t}zZ5ewf=MT3P#N z(??cpk203WJUVPUFrYoNmmMSga10$=_&HA4Qj_Bg+Y^K>xtv(oo+NBp|0fr=rwCit z@u`LFX~LG=PA_cF5Vow<((jqVmi1TG;Vfax{&98@!#Tp1nw?wNo+qv=v1M&^Nb^nJ z=fxPAaV(uL*izwRj9~mO;(`jhFj(HB>35N_yiPB!__?IQF0HW31fv$02g|X+t}wQk z<&M5m_#uz0D(q@82K}xvmh->1@Z;BET#J?8t3BM9KfGO6^mG4=G0$6W(SCj4`H)lj zcfxMj#eZ14zVPXPjZ2dXbY+du3^OX{Rz_@E(bG~xo_ZxT3YtH*5exGsIyykqB z#P2okn%A7Kn)p4&UGtjr)lB=to9n)9|7`tvx7oYqHRr3F_J=puYqb5d_2*q?@0!<~ zuW8yJ-dwL`&;Ib{dTo35hd0;#?AagQT(4u#{_y5{U3>P2H`nXgvp>AKUf-Vm;mvh^ zpLVCk=bG1?Z)n;d-dt~F&;Ib{dSiR`hd0-owEeU7=N;zHHLp3})U-dmx!%m4{o&2^ z=JxClZ?5~>vp>AK9$?S@@aEdDWBuXHwI7%M@aEdrv_HJL_H&Luyt(#s1?LObe!Td@ zn`>Y9tbN!1UeO=kT>EEsd~WAD?;EcDvui%TcAfVJ*M1yvEV=e;c7J$t?Vs`b!<%bA zmvBCDy{l<|cypciG}nIK^tk#yQuKYS==-SX%N##m^nEh=!k>t~llbt_i0P@K z?@Q4a{-SA)J@5L0X^tf{$ByggjdSd{e$KSdKkZr5%%#8Pnz?lSjA@@w+S8`TRkN3k z<$AwrI@kLZ(U*1iTG987=nH>c^yL_Nv*`PF^o74A`m)~MDf+fVU--L4U(bhme9u@n z)7XAsy1Qw{eLwnkwAZ}#vIq0{k+I&U@%3TYc5p*$OWFh5j|)FNj8T_QjrB~npA@!> z8}t3Z+vmoXuxI>!R@k}?YvT)Jc`m;U+lk!BT3&S? z*P!dX2ER6)*Wg#C^YxA^>iC0T)au6y`$@3=?$}z*=g*=aoPCa<^&ZsVTDa37{^Oz z!B`s;SJ)&KHfgYqo=P-)`kKz`Z!%#^jr|@0xyI_A0PK~Mi+IRU{f#~~RreBTj+mzq zK4DW<*i;oZb%jk+Vbco69ypy~ti1_ZJekJh1U6xXO;lm(-T}liy@&-iL$G|U1?yzv z<@3diX6thgpGo*_Htua^!N_s+Gj7qCf3v@2Z5Gk5xs^0E&01lz1Sd;G(P?2Gv?!?=NEp6Z2{p2 zwqW={FC_dhuU#u_VKD~%7BSW}<*=ym!}yC;*y3Ug`gJpw=cBu@WsmD&EXQ-)OCDVN z|NrI++n$B(i51&kh3)AT+g0CajNvR{%l_TF@O!SXWlvtBupJ|8x%OSMu)RpwvKNfB z&#Q9Ht_#}^UaV;S<9R9bTOR+sS6*5$##lx$&Kb)JMtznOjPua)g7H3f1;Ka^yP{wm zYbyzc&y@uumOjC_9``Dd-{y859m@ z2mM#6#JQ?q#JO69^{uegE9`^EJm|*r1>=2KVIK)bt=ACa!Pcy>wFF}yTe)T2ajq@; zkwd=-@7!2lj>iqe7_beCF*d4V1$vJLE;fG^rr^3(9!j_usQrM0Zw&b#FVY{2KW&Q77*zO^0S;wOa z+tI?7-1aPN_Y$_O)za_Y!j|<{)?puE%l@%%5yO7MmYVHf*sj~+i>7hU!j`p>=bO9_ zh%qu_9Umwd{SFeb!45X&=T7t?6+eeo{2W&JIb7IM=Oe-vePr0?yol#0;Rkkf_(2~d z{IDk+TVcml*zpy1La>fLozWQYL}U4yDLgLRC0_!B;TE#`G|iLj*>mxe9+varqThV#?qf{~-T=LI!Y_vTXYb)$J!HDJh3cDd#&YymMY~=h8 zjQrsoD|wEo&5dGA*i997vxtX&KF2wpql@`aYmWD6w-jUg?~rI0wD`XoKbL6sm#lr; z`u~r#|9>Gjw)m_1e0ggzhFfwzxUFD497bb&>ev%t&n5HYkybK)ooj;Ii}9Qh%N>Go zjNK_1`@vm;aopZ57}s9+2*$P6y@GKrzE3cG-Y*!jJP>RmPdXa;jcsM^@BbeSO>G_w zP0kO6=Gc8CH0$c&(7gYDuz3IfQ0wOd2kbw7#L?NFo8#Zx+}PAKaiJ$GwA=c_o7<1I z=Vlkywl-a^#V-DmHMLIrx1pxTT4Qcp!?4E%qc)?SbO#6Q?>Tt$-=9*a7pz}5)70g8 z)7_JP&a|)j!S?X>tZ858=x0p#GEJSIHtoLyMn7eG3DeBulcxO`CblO`du(1mZyp=A z?;3N>Snq_U@7tjn>#fk_@}}vW+Z(2HZm*lpxxHpO=k}`UoZBm=b8au2&bhs0I_LJH z>2cNlQ)4-{PfX|7J~o|W`^a>T?L*T!whv6_*xomtV|&kZj;+P?xMKUxSdQ&m(>b

L9ZcugnoQ@|{)yN+c!<%O<9Q;p^|%%5TwxOj>ll9giMjH7>E!o~541Sd z=kwc*<#XsH;aBY&4?WtrMji~!xMM@JMji;w8oA$eUL(`kSa}|%Hl62TD${u$rZk=B zVG7fE9ws-P=V3C_c^)P;J+3t}tFav0ET(g8Gn>w_&15>qHlyhr+YF|2Y}1?0u}x<> z$2P5Lk4@Em0b}{toIf;mpD#3ZpEophpT~60ZEn*!x4BH`+`5>~xy@-h=hkdG=QfAw zoZIZCb8fSl9#`GF8OyOPZaT-dnCTqbqNa0fi5Z6VV+wgpX(>)5!f zrSaJ4Ua`Hiuzkd|`<>9AqH+H1(dwuD{JUgm=CfyLt|fYfW`36l&HVN@o!^%()yf<9 z6yBpOU17@v%XvH5PrU5GHJ%T! zUAM4(x?;OtVf$jmcJ((KV|Yc_vfkG({Jt)1`Rrta!uBm;JId3wwqarWuCV18oW=eo zH`nZou+8Ub;@!yndOb56=`UH^xWYCOjAMG!VEI^J|9Jg{Hs8$IX2Q?dHj}2!1?wMd z3&F@~OTid#E5SyD&#eWc{@VyfoZAXUoc#qG6XOjKjPXXl;zo`0!u*c3^;oho{*twR z&-u^BejM8a#dzefonX{ukYKFg!Ghs)h+z0s_X}i>hKhc$VHLK$V65eNF8FU~f6}Itv{~eF@ImMO#WgH(n7Cy<_d+by&jyGpJ7tCM7=PrV=K1T|M-(4$gw+h=` zFwXyb2*$b|RbitmY|jeYOEBv4<9S}C#$#oOzhrH1(GRwdV4TbL6^wrS2}ZyD1!L|H z5R4oS6pVE?vsi<~a**gpEC&m=OIul*4iSughgR5k=emQ&T9e;lq96GkE*Rq-AsG1` zSz$+2*wKQqXB{IL;~gs)aULhwfT-v3g6$sc1i^^qOR;w8r}h_OIZ^mz9!?UBekTjY zI#c)T!{;fYA3jg5Fm+6`22ZQ>J6$khIYTgFIa4tDoh2B5hjzAL_&i52`kgBn=cDrk z8yIWie8K29MlkwaP+=De#`$5vF+Q+9ZZ8u3sM*DWajd9wD*e>)PCs={CHG6jcs?;624Z(b`H$I5kr(eHY}TBqA;bAw>4 z33c8a81=kS^y7NsrV6{c!fvUsTPy4~!N}ou!3IVScL>JyuDWhzAG}laV}0ILVRs8g z&F%@-F=F#4*7X`Mw!O%)vzhR-bM(7cFs{?@6O20DFBthgAQ-WX6^wJ&gM#tC=OMv} z^I^e=^AW+uMDC9Y#xd|o<-OFM;=L5HJSKdS-{XR@ex48vpHB+LT<`gwC)2nO!}b!4 z{GJj%iRI}EQ$E>;pAr3F&k9DI&sCV;=cUYV4}Zy;y5|e*`S97P&9wDrbNu^-R@VOe z{+B{?Y`qwLJEXp_mjxRgcCS>}tAgQEjmPgXl}~c4-RDHzAB8~|2!2#-(chr`m@tBb`R*Jn@lfin(M6bOs|;uKTo&cW5-|q+w{tb|7E&Q(tny>CFws*ubTAl zrdLb)H`9HS{?+uFN&jMct)zc8y>`++neLbLkEYj2`UlhNCjGtX^^^Y2^ae?PYkI?^ zzcIaW(r=pH)HM5sUr%qA_}9iaPx>p<{geLE^nj$lFg-Bo&rLJO{_EntjCeN2Ynt?)pR1?7yx*Fg?Qm5$%1`J0k4)J=5$H{&I_H){5(QO|#Bj zzhj#BgRb8;&Hn59Ez|5*uDzaK%UsVr9YlM>_*jHJcd+$o|NYxuH~#Fus;TKU)9)wU zvDLWooW}9+vpBw}bAPS7;PeB zIqnGyKY6WpF8pLRiLtycCocT>aLoUt!TkT1q`@XL)|D6Sd0Nx`lAgwN-hW)@{by>^ zdH$&SN^~ z=sM>(x9ObYT&8o5T}729cob!^VbIwbc&N=rs zopW}bbM9q2=iJluxN=@5Sk4*W^kSXy6K#AU(;pI>>29>%Q>?Lt=r1luNBt|+nh7D>l@2C`(B%K z-qdu?*>%o&6Vo~8jZNpAH!_`bcAazH&~(mu1Jh;BT+?h8Ea$;kn;Xk{^f#UJ*w%E; z!*$MM8`C+DtxehnUW> z4mO=*9b`Jk>N>}|o#`CwK+_(pT1(p-%dzedwmC1((<2MpT}|h_b}^mva-H+q*>uip zC(}8v9Zl!FT<5$-n9g|(H(kyF=drzv<-GPbonzh8bdJ?^j&-!@9P234Io3T)=U83m zSa&y_W8KYk87uzwF_vRJ*mREdAk#Tk*E!Y$P3Kq-FpX~?b$`=2R@XV!{Y>Xr_cdL{ z%DJ+v>rrN#V?EMzj@5OJ^$62B*27KbSPwIuV|ATlJ=AoL^$^o#tepSK+id$9lZ!9P4qWbF9aj&at}Au^wYO$9lBsGFGnh&oGwPz?r6Vtf!mKvAWK&o@P47 zdaCIh>nWyltgdscC!5Z(o@Ba=m2=Ac#&WEFFV`IF`Nngsu5+yCna;7EYdXhzj_Dk$ z>m2LZrgN-knf6%odn;n)+?W|-jVbzZ?sab$6fBQdvhBP>w36lP^^)Zpl&lX9?YU%* zsKmIS%eJ-Oc=??=;;#={rpKFwJzPZKivf#{aFRdnbL1 z=_O1P|IMcJ-?Q9gx?Bg_`AgQuR@m;Zx<%vqp;6&UhuRZ&XHU(V&pY$@ob}7|^sMPTPtTan^YpaoJWo%V&hzx7={!$Qn9lR`xamAk zkD2y*Pv%e7@|t|xm}}O^TcN4Jo1s}FZ-gfA*G=c#UNfC@d)0K#?G@8Gx0g-l++H%B zb9>Qr&g})$Ik)Fcdv2=k9~;ZDePlYv_Mzz<+Xtp|Z10=SvAt(H$JSyx$M&x29NRml z#}(T*#&T?5o6fO)Wje?9rRf~o7p8M;pPSCHeP%kx_NnO{+b5>SwI=>Fme<4|p;;5Z zhh|Ou7MeBjtLdEEFQ#*DKby|E{bV}l_M_>X+YhF5Zr_{ExqW9k=k~4Xan*hNj~Zhe z&vcHhlj$5=N7FgB4yJQ#O{Q~f{}i$PUBvcR#Fnpxnmt@=6PPcLTd@f%Y@%Rl59`#@ zxQEST=QbZB+cx`4)}}P(`yH)w_*MJHD~~mqHlyhr+YF|2Y}1?0 zu}x<>$2P6$9NRReb8J(a_SjV2=QEa%&3Qvp_jy87_qjt;_qj~x+`5>~xy@-h=hkdG z=QfAwoZIZCb8fSl&biHMI_EZv>2cM4F=IKlMNQ|}7BQV;TiA4tt*hxA+d`&uYzvys zu`OUa$2Py|aUC1e6vxKm72Bx`+u2RK-+Wy2IaRmtqgeL}>k+JzyN-s>p2qTf!``O- zxQ6!hpRoK!Mwv1rJvTTJd7cA#bzs0QHIsfI&*7Jw2ApDMu{8y~7 zl`3py!L|#ZeFS4JR;jR6D{Qr3Ie&cmn#%dFZaU9@U*VVi!|#Dyu2nuCSwr+29Xa0k z>VNy~9&Am~k9w~q7}w3pCvmPV`VEco`Uy5D*g6ICy_B{}ah&>IL@SxUj?Z-qpP8*! zFyEhO$|rSMzv$hoH^;x1Y-#&Pm(Uwk{A?^3>-jA^cb97f_I8Dx z_WXbQ4T@Mc5#zx&t+34mBbF~Kv3ylw6TdL-IJYR)h+7ipmV&WfHy3dZjoi1WFy)iu zaZAw;wv}LmV!W*@Y`Dl_kLWj|!nO&YlZ355HOIfVZDsBMU*6$Qv;nS3`Wx>Tw*EB4 z^!kO~AnC!zH%xkvY3l4RZ)ck02+cXn^+4l{?RtReay`R`3dWu{tirYz44lk(NC;IJUEMLp+Tk*4Bh3#Kq2METz92hLe20O@@Ust0K z7JkU%kP16gj6uJ{jOF|f7k(K3hzdKh!j7u2qbux~V5&C98q2lW=VhPO#&Zel<~ZR8 zc6|6jpAfcr|A3t+7_~X6!cMNRQ!4CK!HDIw3OhYmt{eND-=8 z{=A^YpW62b|Ll)eGXGqURx*E`bKtebc+Lj-4{O&IEPocjy@T0jy0vv_{aG@95ueu= zJ~O+aV7*$swf#^&xh}e~=$C!oBp7pavtZ2oEfsdFU>t+D1^{L5!|$P-*Pbib{i5F>H?a1AVDuYXVGjyM zeh&#ojUN_lR9h%bj|j%RJsK>}8*x5nET3~8uh>3O*gh$22Sq$jg)RDNVM`3pgf05n zitTfS?eoHxTD}mr=of`8wS6gU(Ju>IYW+&sqF)VLK9BW!4vIRwX3X~wY~L{LdmH1t z9(|b`-xsLETgG~u9%K)1Z-y;v-+zCB?c0T)9>&P~U1L3y?K_37?|r=wyP}qJbYzMh>YabQ19}8RJ`=qe_v|{^NVf(qTrDk6gwqFWc=J2b+ z_G@9w{C`u}ek*J_ZoeyRz24*6r~IDZ<6_^7euLcH+7E(J+aD|JC&8%Q&w^2_Uj(Bz zzY0bzehbEN;;|Fw@5b^T_D99`&%*XEVLQm3Tl+g~(fBx$I(*RfyI$sI>S9cv8Q6qoJFeqq8qtqDrWFjI(^c5? zf-$!<2*zB_C>V3M?vww`f#YVT=-1gzta+UnuZuCCSFWGtH0|?>ZZ_?6FxVd6<}mGZ zfS%p7&-Jbjt<7fI=Ndh$X}`v&X0w>izjL42wAWt!zYdEU^H>MFiM2&Slh?wb8LMk( zVqeH~&TT={IkyE&=iKHuopYPdbk1#F(>b?!Oy}I@Hl1^u%e3dF>b``r99wVGIksM= zb8J0L=h%9f&arhjonz}}I>)xS=^Wc)rpvvl*P%jfj13)_57?^>MGw=HbV?Z0cMzp?zjWI)AsU}3wRuqDqyh3#Nr%i0-I z*bc4O4l8W87q--Fhr)Kauw@QM6t+7GTh_`>h3(G5mg9Dp!q)3uo+Fz5C2M!LG@eg} zJngSG?pxGsWcXnnxg~5@!H8!!F+Mfeqm{M4|K45LGQXq3mN7>QTi9M<%e?F9?ls9XbCa&DQgWA0_;b^a5Hty26gBuww<|_bSH;##|j= zVJB4BiNSLI`1CcE^FP^ip8u1CUyd#Reu{Ff@_Xb{M8DB)Ztd&W{Ac4D-`!ubcB<${ zy-yR2W31+LwpZM9<9Hm`rwgC3GX$e%kRm4dOyUR7aN3x>~Y1Y`fXRxt9ruEMUb zup0zp9&W5K^<4uT?>C8lu$wFF7Qx8x)(X2#FwPOT3x>}-1f$<&;@E`UDf+Ru-X$30 zsoe2-x9G?Ia*ts2yH_yI0rv?;oc9Zc&j$n}hp`p*V1+#-82iM-f-&ABf)VGVg0U|> zCKx^+7mW3()*0hHQR!D(gHKlaJtY|P{YHsQG zyy(Zd@r4R|Q84z_mjuJ->EgJBy)621JiJn2uL{Q6eN8a-&({T`E^5Dqy&?Jyh;z=H zg3<3S!N~n>!RYsnVC*693Wm?xp4B4yaUOe5Flwgm{l81({=VqP_1Oo4aZU80V65Gb z1S3E7T^Q^kAB%p>(I^ecE5WGs*MgD5 zH-b^mZw2F8SlzRqT7M_{5$8?f+LU?!Ui4#~{U8`V?-D+#=Z~V_zxAwK@BUQjr|uz7 z&3+dB_>4sLgZ(1a{9rp^|MC1O8teLxR-f&kHU1u&djA@l*nbl->={1) z3^oCS*mHCI`!8dilm5rF|3CC_*OL|6S>|Jt{@d)!_e;caY~xsji~P#U|yKdPBDf*m9_uh)S7tIJ>Gf0(Lvq~k3P#PR6O25k7mV7?AQ(Pp z6pVal5{w*XE*Sfv*I<@{`D+{a4{Ng)41dmMD_Hh9dxgzWFs=#QXLG^)wLbpC+MI%| z+=iOE6m0!?-Jj+v7<-|!xr23d5pBhIp4S+2;xElp_+bq=o3CL08hsb|(3L*c7cUrl zocEZ&@R@60GLBc(udlxtt?aiUz@AG+-QA>Q{u<|o$%-7Dk;ei>4!rI*3l=QrSF)U6 z$#Q-r%lVZo=U1|vLpevzhWHO_vx?kd3l;gfJ!_(C!E#*|E?BNh$#Pvvmg`coT$hsN zx|A&Eu)fHhV{6_DTO{(!$6T{Jwzg<1Yu~dMD{A33T$?Xmuw27#1e!;*Op#M8Zq(;2^;QR@!2xu`Eb%bMl7^eD!okB?ijT$hsNx|A%}rDVA-CG!{t z`wwd>ch**UKT5af{J36Gmz;agO711gxtA>GUb38f$#U)`%Q>jrdB5GO$ie%ur*QA~ z{N9GT^r+OOci}UyiIU~Ilq}b!WVtRS%XKN4=g0Bhr^vw>`7Kf8;EeNAnY*)Q|6#3d zziB)-)32O|yq1+J;S*}^ha?MJXYgV#cvyypUTxY6%qS+0sEmh=@^IN)LIlq$S z{7RPdD_PF3WI4Z*dF;%Q%8xa%*N1M|KDRl)Wh(iVEaz9UoL|XuekIHKmCU1N&HKHX za}KV!?p#JN_JZvTmUCaWl6%Rv^8urkEazUboO{Xo;Lx5+=7=1YE8=vm=R;Te*j%ZSL&UFkCCBKs8 z{7RPdD_PF3WI4Z*dDI;5n-@7`wrY`sv)%lMwbd%DuVBrsy_%XLY|RQ=tHRc< zuzrHEhpZzQ^K_Z08I1d==j&vc-$y;K7gsQKPjzBhSHuEaFIfJZA9k^g*VPA!w!Yae zX8(q7Ap8=~&+mGt1NP5;xo+6d`enZxhhO+c;aBzDr0Ba@(RWkP_usWz^xZQ0!nY88 ziD|2%?>5mFzICPVwng6oMc@9SFW1Nei@t-RFMK=Em%V3j(RXO{g%1&Z`G4?+6@7Pz zzVPjfzP>J~<8Wi$Ok=yF>F%Z(cSQ8%GnQV)sN>GYdM95yg)PU!lF4?L!cPxl)MZy= zJ(KOo!Z!b1!0yJDNVdBbwub&)^LrS}Yj;%G5}(^rhtYyHyMeVm1>@M+v68)N6iTyv~6}C^Xd`!{L_rILeex`jsY5NMl%$MJrzKj*NzvxG84iJof2UgfY zf-&a@3&xr|L@?Ipp@Ol;9TtqW>e&+K;l{j9=p!n&M;5k630rD&bl9Sg5w?tZY~lB~ zitX`*?Fqt`TAmoT=#zvkwLLj((WeMoYJFBwPxSv0BbADl)*UlJY`Ix()VtZj>dy%kZ ztz2B#ULtIX@6y8dvWo5Hh3yr>mbG$aVSAOZWvyIY*j^)SSu58Tw$}+;j@#=CTd()H z_9_1^g>kVPL_cbMqhQqbrV6`RFlu*;VASeX!Klq`f>DdxgK?aAw#0ddvAl=fS+TvV zu)SN@Qk#3i7JaXZRZ*Q6QxkkTf+OJ8e*&C+)oJDM} zoA%oK|9j@mW5f1qW3CzNtI+iQGBjg-5t`UPH=T3)%yiD}Q`0%OPfX|BJ~o|m`^a?8 z?L*T!w+~F`+}=0sxv9GUVl2n@v*{e$Po{HhKbp?5{a`xB_PyyG+jpjOY~PyBv3+CO zW8-*r)$Bj4P4#}`b@TKs{%WHw>-<;qllR@@TJLGCZGS6lPZYMy{qJFmK1J9P!ykp; z(<`=r7Pe;zTaJyPc1(G0=yQcF$L3$*7d=MUa;*Lxw&;t(mi^GHO>Q%PN|CpbW zQLm;?|I6^%p~5-}M*f`y8|uSY8&5Fu8(%Pv;R%A}bx52O8p~^IqKa+j!ggX|OKm0z zTlA#DmKY{0{7zo6ouaUvQrJ?u8I0v~ z`b?&M4;mWd&KP}}n|w~6#aM5%qlPnwEo(TxFPXLQ#ky5k_X_JF zVq?6X#_}HDOV~1>J#GEv{iv7ea{tA)w}_{|n_F9=!j`PCr7CP`!H8v<3R^Z9$GFGG z@#fzj}psSFW%=6}C#SoG(6o zE#`Hzs<5Rt-HO`ujXLHvJ~Z-JE&OCgj_P}WsHytCBCdn`3ZJmm3+DSaO?@v9Hj z^FHhiYZlD+ds^w!_kLQ*yuLKp1iWC+&GGLYo^SlymbGpCi4%A2?=P@-?f=etY|?8L z@f&LD;ywd#=UiFM;52RCUByB)=RbNp{F80$oRCk}JUTHqSchZ+#$F>n8y zc_N=!8)z)QC)iH-A+|vkHn_ru2*&t_i8?h$ zj6+2~*su!QzQT5>u;CRpLNM0jjup03uv~lk?QG1?=jdH3entvL4b}Is!1x}PyywFF zdsxb20;ax`g*bN=vA}in#K77ZEf{%? zc=^9M({E4F54Kl@?Hw%VK|fz_Igfo!=WDoqgkNf~Uxn=-EZ3QS2gDfY11o;qwp;_u z*Fl2e^WX|Qq{0rZu)_o+mcuLT2*HT+$O=2E!j2Y<^T{y<%kK}=cXbT%5?edA=;yBy z=Wzw=;V@e1GrvD5SzcQub9Y0+=kdjO{xY$gAQ(BEC>V2el3<)WP8N*oqEiGL9dVv2 z*zUnj6AYiH3q}rSRM?q+z*$?dU!p{}L z55H?ueqg1aD}^8CPWgeAey$RJ_zcCrzoe`W%=_i}x?1>QO-DD zuCTpc*b?6jh3$>PmKbg-Y;P8}tczO;+gpV#bAMZ5d%LhDw>t{kJB2N4q4ax~u;myp z>u|TQ<=DBWh~eI_&3gy&`genr=XGkYz9)n+?h`)g=ikv%`lRmni+;?Xe_u=4k2&x) zQ!>W$?`|m>?{O*nQSS%DcIYN(IF7~>Il#sjj5&HnK4Gs2Mh>qEMh>cG^qWb{ z5pljIe6l897mQk~@mLdG#CWhbDn8#7jDBxb*xQ1!hrA=$P(NI)y(<`dYl~p4iT8r# z>mA~J-mY`+k;)bh))MSmr1hqhMM z?#~^z=x>BAwf;72(cgtFpNDuo$@_a_-ArTqqv`zrSNsrtnVVi|E`B!FJNf!4Y&mE7 z`*3RbOX0`gPm}j=#(E~(Ukls(|5tRgV_}J8`+H$)=-)B_FqW^m|E$>lRoMP5Y^m8l zg>BPkK5FCpBJ%EFEXUBXV%w>(9Z%R&v+)bt354xX56jwwh3!Pbmig~o*iI~LIc_H@ zY`xy)c^lhF!w+l2E!m?d6O6rfd9nAgM=v72M{BSeIu9euZriEZ1gu^z-kC@_f;JCsej&j1!95Y$(QLeQi`>8&}vSf^qz8Dj0LTS%qz0 zVOs>t`Qp>pVqP~}3fo~($9;<0oFw)n_PMRXPhL0VsJ`oonyT-l8XMzpEqubZ5sVzS z6^#EstiNEa?`dE6fg6u+*3baak9|Yk!=624V5OhB7d^3TC;CydL4pm8nhh3=+=mFp z9F4YPrpz6-XN8?97;`jKj0YQ5VcQEv?mGy!OXNPh!bS+jdf8F1Vc~Np!N_4}!8i|} z5%>M?7=HX|K>T}`)_zS*_!%kuunu+=jCHrI_-+B_dW2w%x0~?Eb;s_4(QgmIhDA>5 zyKd+=O7vqsM+-*oLq!gZr@l9detQa^^xI1?t~>S?44=nVV!2-Ii{!qK@Cn;jumMre z{RCre?JpQU4-jlv_*CDU#J-`vUxwrVK;e^P=b#Ex^T)N=!J;3vK149aQ}X~jRPuJF{MxFeL zy0jg)^}U_{$6oC>%;)T5+-qYEnX}VHY_Kx~Bi}Ox;}}0nF!DWHFy>`M<(lan(T_OK z6^!+Io?!GlUoc`BBN+SJ1%i>^g~2*`$)h#L^9!xlH|_s#{7a#kL%y@710C&|JowI+ zj)@J6e_v!{=J#Y5M-1a7Tdu7yZDsBM5B(+47k*i2=JN9JmB;0HxFT%PR|?x<5yMqs zi@rK+bFR%{driR>NqKB}yNmXsssF{u#&3 zq49ftXpZ?CLUYY{Q)t%cjiEW;-C{bwKfEvo)V1X?PYx92MKO2uZJ`a7a z=*O|)J)SR^zs9`3AQ-iNQ832b3pu?8vk!<%ABjZ5$ijmZmfrQ1;eMhk2CYsBKnczBBD09Q2Yqojthqvk9@t=+R0LR{HkzZ#wvF86@413a7=Chk=j)N~vcTf5Y z(>+X+^XH~}CjFV|UZ%N*_|$aoq(3peglY2o*z}U7sqaUoJ#MbmJT`29H|CnLe+y0D zUqdtYFQLihXVW>ipG@c6el(qP`@wY1?R(QXx9?2n+`cuPbNj}0&h2Z{o?B;svX<-K z>GQ_eI-1V0bugV{Ycic<`^S9e*#0)1WBbc=j_ps=IkrDck1MuGi`XV9Vw<>#t#c9E zM5c3W6PnJkO<+35HooZ`+jyqOwI*gVme<6Lp;;3%gl0`lADT5Wo$0(LrZt`GK8@*I z_o+?ix=&>~=QgG3oZA$pb8eHH&bduydR%qyVl2ltr|BG9v*{e$9Hw(@vzyMb&1O2s zHmm6z+bpJYY%`no*zz@cv%h3*t_qvG!sZFawXgePuKeER{(b_zLF`l6w%I#aTf>;! z(&iOo;%C0fnDZO!iW7TYK=>IJd{o4f>oFwQf?^EVLd6(eD`PAy*d8&)B4P~KqQw}C z2|w(kiwib7#^@%-fORj%=pp=2tDeFS`*yEZ*8bfqdqiFHJY!pq*}F345@O7ekr!j` z9(u`&Z8_#rVoc7LOA9~LVA)pI{$602itTd7@;Y2TZ1Y;<{Jnzk!`!ae%G%H8D^+Y) zHkQZi6SjHGQBki|gddKtRSQ3>30rE?x3FDZ*pkZ{h3%Tcmi50@VY{}lWgYh`Y}XOC z-yKq4B$RXwJuLh30&`rs;gn+t$X)=e%u9=X2iHrt>*(E7SR$x25TP z&fCItKId(2I-m14Gd-^J@i1dKwxOnTY(q@v*an-!0{Oy}74HJxMI$8?TuZ__!py-erW_B5Si8*MtrHp;Zemd{Vk{*tw| z?fA;qAqP}!*D7rH702d~Hn&X&hM#==?GfxC!A1u=STOn>CdS)6`W+(rjSP0EV4VNd zc5m7pSKIfJlr7|wZ2m@V!2B&YOVUQ=2e`mv%7^)*gb-g!@Yu$ zgQ^+*`2LK%28r`N;gdCSzhKl_jmI4McXEtty*yCyIaV-yK3HK73C14suwd-Fj|j%z z`lw*6iN}KF>mA~J+?c<2KtEBjeX_89O4w5Gr^6QgjId?QXA8g2RcxOxY+n$z)bho! zMZY9$sqM>Qi+)AeQtMa47X4b-^7}}yCwaeatea_U-!$#-8yV+~=*!&nN^|kHvEIqo zTVcy_;qMQr;X8#Ne=kbjEyj8#+jk3Ff6qSL9^UvK6a2nc*t!kuePj8W`-6(@hlTA& z!j_tST-bghY>DC1!uGR@?dOH<7s8gBeOcIkC2X0)uM696ge~*`ZDISJu;sY@zOeOr zm-`g9KM2Mi{i9&)wf@~GrBC+g{!jgvaV}AeIq~mtDSa}x!v!PGpDJs0vpf6TpgoR-xS{=e9$*qx8qib{jnfh`7#9blm% zAdTJK-C}olcQi$LsEOj#}o()JZ#y&=TLY zN!w0niD9~=onB~}7c(U7j6%!U&y=(?3oW_LlC-l5Eps9JolR(2 zd6ND&c*85w`ukkeVxdYO3s<;BD%_$KZZW|T%i2oHIZzti0pO2}yS<9W1T~D7TrE!)E`ZKZ_qL~4#Hm7FKcUx@ZO%~>n;{c+wi_?r%srk+xpwiHxmC~!0VO1rU~z%-~Qv4 zH`^EP@F!co7+dfrX4|KH?Jxg-I2(fTsSByp~7IR3nSEyF%8>{A{a+HH(;Oy64<%pBUPVEWp!U~<{Q zaLjFU!!ftb49DF18IHMaYB=V$iQ$-AU&ArCjSa`#HZttFsk-lMT#Rie!!foU4aeAa zFdSnWU^vFs-*Aj=d&4od?F`4*wlzGl*oGJvV;gKZ#B_AxHbiM`fbjBS6zF}D2-$Jq8Y z?6JkYr{Ql|JFLPTUg3@q9QUIm1=ncqt*Pm#3U_pcJEm}akJ#NYH}i;f2Ae%U5nq6O4D)0uiE40c&7Dyp1olDovmQz^Q;9kpJy=~=geu=SDZ7a z8jf@36vJ`OoNPGGnUf61Idh`nIA=~U9OumOh9@?k&o?f{cAnuF+qs5gZ08t`v7K!= z#&(wB7~7eKV{B&_jmxm{*B=60#!nA;_W zV{R84j=5cAIOcYt;h5V6hGTBy3{R}?w-^^=yV-Dz?Iy!9wi^w{*lsW!W4qpPjO{wZ zF}79pd*vYjf&Km?bsPw1KhGC5SoT2O!?pYBhu9!5CY%%G{)&i&IA0YUcB((>QsZ8$_?%kANq*`K6k<_luCN|n7yYpp zydgN&`r#R-sJH=;w28VmsxcFTAp6H)Cy)XR0 zeUSWoDEv^H@5)+=@xXnQ{Cq6@Q2$Q~7yZC}n*4kw{IJI=KXBR4=fV$Tr~JTWKVJwx z-0!{;^?`G{I9^{0Kg@}*3K#1G_jU5)b@zKh)G|lDN!o9PmiWF)+V6#y7=B3FABC2A z@l(?NEVPXMFG>5W(30D4N&CCdG8eMnKZKSwp6l?Z(6V;^N-_MeXye{NywiL>;XY06 z)%hv(@wf1aou6Hkp9iS>Kf;dj^K)#n9b@2gCUf-HM&v+$wfrU(cGSD+s|ovK-<(|7 zu`b&PyMgY?+LXc$pDMr6EoMz^g&j3hc5q8SF3Mu~GW!BLkP zDmkb#XXtP4irqYdV_wf9IC57$;gnD2{OrOfakdv6b8t?<@w_v+s2SW8f@3|XIN4Ve z$2cf%*D`+di#VyLsu`T}$()!{}DbD0nRk~x2mSlT6@k((}Y-JL8=`6T!0lbxH;-wcA=wfLN| z!p$T&=F7~5i_corbJk|=zyG^1i_j9=Y(+~vvkNWUoJGt1YYw4pl)UCD+V%y1DgGWa zxy)VssArXVD%`vkZoUdPf8k>7h^1j%{JGptH+Ta7%`r#hzd)sr_7!fy3b#;&TUcKh6Mpzv!r~QfPVsl^v8$cSv4rr+SS(rLI#jr&3K#Rmr_Y%< zzDo-~?DHz_aXta8EhFs4H8b{iG@th`TF)8!?;(cE+=xp?(#Ri2R#nU3%dthZfym@ zF%ByVZcvL^(@KJ4ets|hPXD-Kx3aKK@l3X+;8_1_RXCO3&{CJRg&nb|aUs8TgdNWp>k1B^>j{qWRA;H| zP~u!)*f9>n#n{qcA7KZ#f#3#}{x%dG{cR*T_V0}a$InOm3Xb`@iQuTqrWNk=%DtqY zuw(o-6CAM|RJqS=F6@YN3&F9*&k|#V-Il^`P#K4<1joAFy25QEIDBqfxVUe?ZD*X{ z|H0cAZJQF0zfhjc;T%aAAIDemccv#|mj?$DJxgVbF+y);r{F}IEt96pZ|TyO8x+VO(J=Lv!%mJvX-u!|D}&Jo4NnS`;1}>e^$Ya-Py%g?3Xb(r)a_F3N106SG3^si#F!k zDDjRH9Cf&0BHWOf4(@fK;Mi|fZOKvny?n;@BH@!*E*2bi$|p5e+^CZKCBi5B*_)rc zsP*SC>gj(ke`4IF!Y6TFCOCXvUg53~9C2PLIC8&AaQvL=YQZsn*Ay zt7xxJ+8czH+T2*Q;G2Y&7;a8}Z>eZ+P1@UpmRjCkwBS30mfGG~wBWmhmRjFkwBUP+ zmftn@dXo3O#&tG~_I|^D-=m-ViY;T~&tcT%LF0M^Uk?;5>)!9XXdg;`x*A8`j~Lf2 zXdg~mzb_85%Nu75;P=s_bq(BO#&z&FN_)JbeIjX}6k2NbRMI{zw8ZdC(mq?!K9{u5 z3oSK!A!)}8Eo1m%(!M0LjQ`6?`-;%AZeLAWulL0EscFUEDWOKM2|H^2y5O*Tqr$x@ zIBNHn;HcHxf}=L?2##93TR7H^t{-#3`lu#fd{ySz~XjA7i%6H+F6ZrOkcAUjP1+? z)7MM|6Z?#Yea?bsFdTE6-f+xqI>Rxyc7|ha(;AMsO=CFbHnriHTU*1Po2vVK#>Lp? zH5_A`$8d~oZo@IQxeUkH<}@5*o5OI7ZFa*kw%H7OY^+yD4gbg5y7rmN_fpzpANg0U z=N4nsS*#tnE`nq1<}dxT2Y3(6^G1bJKKY*g0>Tciz2Jy*!3wvU$YGQRYi;!kw@~rP zUh4XW|6^_8X70avxrpJ0z3-Q|MMXT!;l%_urp2sj@d~$u;PC1F=kaCDDWBxOnaF=w ziDlIar+khqc1wy_7$?q`iF1bjx>x#BbB2AoSH-S_h=ut}e`@{s_%ny5!;}abelQ~j+{DGxXyxOu5}R{^S7(uSQFg@ zM=UBo#-+QkV@)WZ%wJU(_BGWX`%@3mAG!Aw966}*V~#9a?BX0@{CXJ|pO@Eu)(5Zk z{Q&K9NxPoV63_BQ3*JCz@w-CuyKzOkV$yCZwB)^(Jqvkm;LU}WF*ArT5vwqQn`v@(y+93Jeu%g{4X*U*HYT37F!J7yzwcWI6!Tp4mT5ndg z;LVGc=LfGRd2eByKRcq`%CJ8>($AK~ma*w>9Cg{oxE{gR)@AQS#p2 zIDa0+?{-P+&+Xg~#@br&XHm5Mlh&wzxdFy?@Ha}^p`zU}X?GG@YPNII?jp3rH!x{; zt!Q^k+TDegnhi?YJ%pAq9GtX6gqHCinzVZgE$enz(t5q~zTYUd?~^_U9Pm<$mbv1H zc!w81k)!q_1lMq5Ya<26Gu^>r55=cC0}1Z9ik+X8llNX~bGYEhPx*vXKDn>^*+scO z;`DQJGRHmt7{T$pGfLD9ZnWTd7E*EYOrto)L2;~=lSG`0?hm`YMSkRPy0ByX)YxLD&hDYVeJcI!E4YS-Zf!ro5zEygPGV8# z-oV{du~X-*kb|;=Q+CXy{Y4IN2MCV&qWXhV{ZW@YMSiTKm+gL$=LB^*P{cx<2MG>4 z)gN`KaR*m?J|W^HKXt|lv8XduSPzGY{@4o+6&&mRFu`HRnIrL@%DO#V*pb5#g5w!P z^+$e+V@@1deDX}`BMoI3I?alB3wewY)d7cSNZ?u_Ke z>+b76YMCQvChb{5OMGW1?KwhA4Cf~8c|yy)I6rB}2`yuPLDF6*wB&YC(q1gI%!Tat z5}{>{=Q>;}w5*-WQVf?DZQMJE*Uv=B`!uyz=cCZa6~ZTWel|_^N!_m$c8s5&QZ*9nfgsB>KKsm_#vn?(Gr0k}y8$GrCQUvkaJUHODl zKAH1>3!lW&R>Z;_Y!V#LJJ*Yv!QCJ@)`NYs8^4ak6Tf0%@ zfZf!hKjNH8aKt&S;CSv_>@)9G|9o;fVaNF0B;sWJW)OCa(apk+F;Y3upW>*?ETTWg zR@t#f@6&C7978yzJRC;T>A>QV1-lXln~23A{MxN3&;14y_9hCS%2}_?>^J|_m<)B z7k(S=z}f?XBgd7+nJSHDtEQ$0g&o{O74G4}#XPX{@r`*rYS{CjJtF+FejgJY>*;a9 zvClkFxcGbv_oQ*&Klmx(hrR4+!SPJ+jOYXI+0qC2IiaPG=ZhBnLeZ+{!STWm^Wnwf z2mF%oL(VS?KX9)UKj2r3wrv?O|9Js?j^ZkT*zT$He=T@q4gT?w_Z+uVm2lsx3`#^BS z*{2fA1{LlD5hrtKL$O!EeJElfhmQn@-NzN~6TuDf@>}~6CBUgp9_w;{YBy8 z`X ze2>ES5%QkIINwL`+a_t_&osBUwa_zYCrw(T{`E7NaUJ}P(k8EHr%2i2Qb}CxtaNNVXBtLNv>tXPi5JeocNR zVoZAoJM!o$IP8|KaJ>Y_*e)kH#&UVVF?KtSpD>5P{+6{Bie1|l*y?xW*Dx+VkM=g~ zW1U<=xT=J^XRIE zo|~%s*2cxywlW-J+tP51Z41LOw#^O4*fuj9W9w%)#97~3X>3IOegV;h4t`hGQNB497hB8;*HwZ#d?$o#B|rwuUD*e})?u z=g+W$2fK;2Jqu?33@w=XGsJMrZLr~(+a88vZi5WR+;%q{bKA{u%xzc0F}Hz+V{W?` zo><-YF)qfox8WGuUWQ|AV-3gH#u$#VjW!%(8)Z1gHqvm6ZG_>8%^yFv#_z?neJk4C zK5MN#|6{jm{c|?<)cwp)yvO%@ZsOYglXeTCWuG~qXu(?xEioLJ{BBp#9+b2Lgmy?N z(%R!&as%%qwA=>{E`GrSg?4oDdq~lOcQ0D*y)GlS^DFnkL(PwmOSGe2hY1d!hgY~G z1V{cy3XXh_5*+y*Ex5rRy0v2p$1{q*N1VqR7w6V-747j!dxFqXn-hx`e3H--!^z3- zDHZLhNqd^mQp?ke7JP=#Qrk0&7JQb_QtPve7JN?8s^`aZjr08q?fHg%zoMV>iY;T~ zdlPlJz_=cUso}VyWiI-@g!aPZ$M-<;zSub51JPcTwDI}zy7c^bNz$5|fBjr)T>K1t zSw(w!(q1974(ZxC9>@W!OQNoX1Wo0Ikyp&jf+ zwsvdM-X^q-W@S^;9-pOpy!3UK;dsBj({Q}s-eK7DP|wNF8W*2)o+)^6@%?nc^z~H1 zBMSed;h5VKhGTAz8;-d>W;o{dsNtB~BZgyc4;zlTJ!ClM_Ml;D7>=>M zY&gdDlHnNJi-u!t;|<5yUN9VEd){!2?K#8wzSr=#ti4v@UaxR(6i(&)p>Z+a4-ChA z-!~ldea~>r_g%v=-**hheBU-4^L@*3%=b;h6Pst>85if-w*?P&6KmfT%sl(LVCLCZ zhGTAD8jiVrVL0aYx#5`GXNF^LpBj$2ePTG~_OaoZ+ed~cHqU-HF2?qo;TYSmhGT5M z7>=?1Y&gdDli?WKkA`DxKNya&eQ!9=GwSw7;@pn*P*xaa(xVsW2i{U2*jzvKV5 zerAPJXMhpQKf))msI$tb+5L9!$^GH8)}J~9jQ$iy?oHoL*dO^lSIK>sXC}0xKXuj> z`Mp}nPx&OyHX=@ZP9nI5hi7e4!3}Gnn$#Im_*Cav!Ex>p^V8kI`MF1V{HTkcJ(W4~ zQ$FF8PkvA31L2cc{LG`=-za~}8s{Iy9GGX5iJHMpF1Rr*E}K-GeBVuRjDzBMCROJ~ zQBPGfIOVfZ;+#U{fSo$8hB*B^q+Czp^z)B0$ItTBIX;Zh55f*^N)adH=jTJ^{uraF zgdJm~a-cuOQ5QcqD)-0OD!Y*`x2EdCyjFJ1Yd<3@_lI3uksmpHE9@9QHMZDI^ZA7R zQBQR?5Oq-;aZWAbWN&?|vbUo8(W)EaLlEd1jpK( zS#YesSp>&gnN@K3oUL$EwuILFFFDNK%>DN>j(G|u|2YaCQf%iem}l&{3uescGVI^S zLO*Xa-+In9im&-9-24@;QQ;P-aP2GHf`X$43kivy?kC0ryJdvW2TO+?1vjK{odn08(^+tXid`4M z5ocGy?N;o%39f(Px>vX!f@5FlDLB@;ItPR}mlby7N-Vtu$9}Y&;CN1IZ+|B}j~~ww z%L}_5N`ET|Zb0Ey6dW~MNpRze-O7R+UAU`0^esXvuvY!7(S+6&&k* zz>Dsn_4twddctm8>2H0(;j@q6h;swMF|Rii+>p}WMuKAx*;sI_&Ax)eZj%bPY2hX* z{rd}c{fzT{7`$1<&*p+-zuBU~ZCSY3KXzLg=lz4X7Ji6r8{r3T+u{ejo$$lGZ~F?@ zU-W_90OLA%@k-l4_@VzDE8I?^5A1d}E{?}8Ld(87(6|`S8D*c1bw;~u(w<$>?v}I{ zRJ3o~o*84fSZIfodhMS4UM94BjvAD-R|+lr_8v)ljnMME)aKLH@!nCi@p~fdZG+8k z-cNWA7$P|O7%DjS**yg}q||4a;Mhxt3y#kTBLv51kCB37t&I{KK1U0VSjH4?suH)q zkl%I9d;Wi)OKvWh+FW1okYam7!R&E270k2BjRo&j^tTxH^@07^N}PKYE`EM$6hDVO z+0@jweC;oX8=gAhb_q8UUMk?DZ?gZ>6Y!DdFU}pWId$JV#oRO4amEIO9BcS+!Qt}=!Lg=}6ddd3D8W(lqbuAof*Vq5bFARV?>ND& zQtXZw9QTlU??ERCKf9JbP8Qsd!ktpM*azIH#>I7fn$Xh6=|u}ZqiEy& zq8?`oKX7LiKj5>4A8K%p@B?>l@dG}uXs0WA`wR2vndbY;fBWGJ1rx)w1+xyHFPL?B zzWI*7PjI=I;cfw6Y`D8&`n$+*{CVSrhJ8OG{tFEI_y34@ zoZ*S>30DY?F}qT5_*7??(ElZ0P1qmYa&|Azv9PwU68*78uNE9@`x?Pv$9W-f-m#}& zE9_YJ*9mTjhi*;vM}CTOem+gMqh`vEF?vbFLS0@K9JPK$aM*F?OPov8>{Ve$U0xI15KqgR>W};s z$J~8g_=I~yaMb0^3ip=a$lcFp$>RXWSuL?H)bnlO6Yd?su|Fs~IAuqzl^vY2V;%Xq zHaR~y&bW#BF-Gr-9N^v)9Al*Zel7jIFYNI7f#4X24+Tf9KN1|Vd@ML>t?ZcdDo*Cv zC&DM(r-CDg&jd#fs%F@6u1p*w;{06rWKMh`IBKo>V~nQyx-wrXKEDziKEJMT-w1Aq zhhXhn!LjdtCph-j?*+%4_@Qv|d51WEG|oRmf`6)Le@@z8gqC{$TD0KbgqA*kPk#TX zX#Y&wzl3&3b0zJ(xuOOCEwt43pP~gfedkVEKZ{fAHpY3)z>^g1q@`cy$a_-bItT6K zhW)c9{Y+MD8Jq6LQI{!=>k)iSQM9ZJ|4fN?s^rH%zmoUV#&rwYwn^)sjR)D~Z4O%t zJ?$00(7x1HT)Y1-wIaT5B0lcFi;BG1A9{*f5qlrO zk@I50CtObv3-w>TVyALo?kbM{mMA{ibG^uo=1RAmspMZroYx%rE-7N;xu}ESsPD=m zU!IFrdBFo|J->N=QylxCpG}kP$ag)#5$95sIF}Y2aV{e`@>@&f$Mai9VaJ^8BslhC z)gSpOjxp+7e8#bc>tbBIe{~i8Q>Sji4_x=;r-$%EZ8j^liSfYoOn#OXeyD%1!bLxD z%OyX{3qQ;$}N&chk4vr)CbP(;&`nj{4gh0E?le++$zbB*WK4n z)G|j_P1@CjmiSgr+TKD-3~MCqnnKIGSSxAQ7Fx!Boupk?XvuB8q+MTVnG4x(AE9N9 z=Q?a4w5**CQw$pwZM+u~?=Z2iQ2PObqmPY+PwaLScGSJEuw(po7Iut*&zZaz)89a0 zM}M{ab`y5gdlS(g`{p2F$GY59*m3_HEbQ>9^5gk-sIa4E%8vWII+KTSR~&Kn6FKla zIQe@Mjsu)J6NFec6F!MWon1nHmx<>eeAfC?XM)h5;>dk-5hwXgRmuJEXD6%+{i(Ar z$nV}te#$3tZXx2t=azzNczD*f5*%~)u*y7B=UdPp=OD!W1@0u_le!!wIPz0I;gnDI zB|n!U?-RswhTzyII1eG_z&zVp)C_JL!LetlINAFZ$2cgCz3xjADsGI<$Ji?+Roy07_q3osZM@#i~pNPe)QM2rHfX-=d!!y(AnNcyPM(o zdoH^g?&@(B+ku9=1-y&l?uPk2mz@p!_anhO8SZJAICnI>tYLE6!LZMRcK%|mlU-=X z80VP2M;DCksDkNhWWnSz!f?!OxZ#-FFvBsoJq^d)h8m8!4KW;Z8*Dh{wuj-E+aSZ9 zo2vVP#>LnUFdSps-*Aj=Kf^J$eGSLh_Awk|+uLxAZ7;(ywy}mM7TZzA#n_HC9Ai7e zaE$G6!!fqQ49D0GH5_9*#Bhx5V8b!CgA7k>PMmIBoD-)N%$zv2VCKXr1v4j3HXL(1 z$#Bf=M8h$+6AZ`PjyD{0JI-*-?O4Mxw_^;)+>SOpvAT~lF2;7g;TYR_hGT5!8ji7@ zV>rfkw&57tS%zb5XBv*NonhExi+gLs-?DZ=g}boAT~s*rx!KPH=xDwfJO2KE)Hb|x zYkn4kYiSo3ziQvu*Y?mjNA@Y0e)cYyIkH#5%#pE%<2YPxeZ_IO%5WTqD-FkSxWaH8 zhszDeak$KI9EVE{$8osC@WkfGEyl&zZZ;fayUB2j?MA~fwi^t`*seDmW4q39jO|*( zF}7IXj~^_*{>fcn7ZFzFm=DLVCsIa;h5V!hGTAb8;-f%WjN+`r{S2}9fo6W zw;PVR-DWuEcB|ou)%_{sVr)+uje$X9dUe@N*UJc@Z0SFBsR+{Da3={Jbc* zktMd5D%{J3i*>~A72`UZ5AdtP53#)_{J_0l{D9vOe)zvmZ&tXsL?76_ZCrI);8i z%lD@~Nq)BwTAuSiP1>!6c63uq`z&d<6I%Ae%|2@#@9sq#*Bj%#$aweP`rZt;Xob7- ztqJY;4EDLjnthJ)x2%03xG^nM)0cw#*Ut8r?34BMm9QIDe10uB#!B_aX9UGj-*1Y~ zxYpsmH7?HA??nI9>3iV^?uX>(N8yLs_?Z-WO~Cz>{QNBZQ2$>F7xRMqHTn5X_~HJn z{J>>DzY9N%o$>>h{rn;PurK^9>I3&@%Ih!ThdJ@T!o~W){hj>$Becwsrte$Nkv4|& z_z~YENjs^~62oLkJGsy@FQ!P^DTS7?pDJnF3N5)!owU;kEps9JomOaB@y_y1OYI{^?PnDneas+yVmG7UsQXNUWBg_o9An^fChvdrH;drtubQ8* zqu#R#j(u}>!Lcsq5FGcrIR%GLl^^%Mxr7}xQ+AB)+`^9Wn@4ceY+k`J4)Y1_Upp~h zsOS8`j+!+DM_p8Z6YVN3A;uj#!ow z9JN+<%y|_j^K5D16K)y7kwZtpk%Oukb}NW6BF;|2Cv&2+;Hb6gk2$f1=nt+-#b;N+ zVb`s~br&3aNDslWC-)Q_`|h%WV@~ucTznQH&gG2jXmc36d_}uL(yl19)O)3(1+Of$ z^tnp%yJ|(dTGFmAwA8Y9(Sp|yT57vy(Sp|!T57#^(Sp}0TJ?L~>l)YD>Wg-L!~Q)Z z`dP2oGB(}ASZrWi55v^3Ptmf^`1gjW;fBeNe;PQQR6KH$9QZhIM&Tpg=1_y_kH~zYY*AA zbn{m&|7GtZhOH|r~ZPYP6Gr-opul$b=t9Tu};LflW}o> z*}0%veW~G?qGhee{bjG@C+;u%7#H`Ky_2?+^}+r! z%GN^MU-nH}UrmhLe#XV`&+K2(9+0#L3N3jal(Yv6EioLDw1-x-hb8UdLd#q^B598l zTIR}8Nqe-=GFOgC+GB;5b$eXW9xt>F9}sJImY>tLpr)n+Kl86ze=ichClo)-BUkdg zc%tBV?mNWhM)t|`;;Q2(wbKO0JW~CUpW+y!(~D2`1@8y$4CCUuJ5%&eoz4<|;Lc8d&Jlj7jh|EG;}qk8 zJ2&|`Pxzt!=NB&efg6|nTp;`~r<5PK?B_z^hp|(B;If~KgdgUypAVJm19x%C>k{FI zIdN&>VtwE)OMWgFTIR?VNqeQx65mxxd$rIK!!=2JtS%<=0f&+i_o&ha~*CKTGq~ODTdpNHtwm!duy2+k)!sU!xTCCxI_5F?oPo`_qzng z_}wiy#=z%H-izt)9>LLHEx&t(9reCXaO|7+3yyX9fZ(_nKPWhSs{FVwKP2p^nX+SS zA1-!ry*9jz){aQ`7P z7cSO`IG-^teg=8AqJ1uDpBGxj?S-NRj~7~Ecrp2XsiJ*3Xf?pF_YWsT8 zg5MBYYW-%>g5N4y^%>-C<2u{yLi?`a_!;D#V$0ZcH!XE}-?$!zso{G?%YGF-$DYWD{bJDgI zTH>1}X{V`ZXHD96Ld#s4J!z*GTIR}}Njsy^GX8TV?aV?u*b8TE?xf}Ccx_wgrcU-s zyx*AC=h=wDeQh}Y{?%88eNJ$1{?f3|3Gf$&eeMk}_MaQ}xd;Bta8JXG(Wi!e{t@dZ zhJ6oh=P%Ye*@gB`;~dl19|dFkd%^VeTfyY=tKpd2FNR}oKO2s@{bV@i_M_pL+Yg3g zZr>Y@xqW9i=Ju^&&rQ{RvL9N-q)t9 zXs1uw8HASF%viMGnS_>F&7AzsQqj(uw6h5LwVktQ!E*^MwVu0Z!SfWY zy06V^TxXj>Xy-TV^PYa@E4GYHyss@_To2Py!$#3EhvR*%ee&aLiM$swF5cG`OxjMS z<-W%Gg^bO@No&%7*NuNJi1(`6dAR;rAZo8DYmU6#FM0a&f7I5_!}VwWsJ%w`B^Q6r zkJ{RKxc+P(wKoXAjG;fzN9|2Q%lP{a0L&sW|#yLiil(j;-DIp8wnWdz5hMtg-nLVo_(RG57ze{QoQZ zQ*~N0)u}_uDbAIAMpqpF%*FA~oM%W4mrDIPM@~y8u9K6ct?{{k@n4-fL)dHPoWkca z$!Fv`CeHU%n({fUke!m9zfFIg1;@V9rNVU;9P^-C;o^8w&mPU(fA@{_!9I@o>22E1hFOcN8;;L+s~L{ZcdHtX&v&aB zj?Z^18}_+3y!5q_;rQ&iqT%@Lxq@MjO+8m{Y@El+bM;0AGY%UTOkW!mOfG#4cQ$+Q z`i5g}>lu!@t!p^uwvOSL+uDX>ZfhBixvgn9=C+36iPe28<2r}fwlo}L+rn^+ZF9ph zw#^L3*!mfcv2AKN#V{W?}j=Ak-IOevi;h5V% z!!fsA49DDdHXL)?$?(MLzL#;GLu_LW$JoXgj8aExt);TYR+!!fpD zhCQ};k8b!|*7mM&`&77n3m5lF#>&s!^Jfq6V&8gCt)C|v{+6}*jdLw+zv5Tz8;^`{ zJx3lcn0_8Am^t!b!OW2d40pDUz=!?NdK?Zl9LM1h!*LuAHXO&{Aj5GS4m2Fc;Q+&N z9QHRnu{m`-ljO|FnF}5QN$Jh=x?6Il3pJQBH zn`akH-OnnRx}RAvbw0yz=aAd!hGTB08IHM~YB=V0is6{s$%bQYCmD{pooG1bc7oxF z)%_CVI)~UUHXLKS$Z(A9Lc=k(3k=8D#u<*Woo_hCcAnuF+qs4(wl>sqYcss=|F+%_#`#;;E-QURyJ0O<)8&HW`DmwCT-4g-{ibuuheYhvn};(G;@zX_?6`w*V-4Y@Z%4D zX?W@2=krt_rx;gt-rf$y_G@8B-M))g&$)5OZb8N zU-1L}Tlis}{ZrwZe)M7=u>aJh=Kg|R8{;~J_op)4C*Ua!Z=7)7fTu7%uGz^AvyZ!M zGQ-1y-tq8&CpFC6^S384obOTO-?qX{UAQ}@%G%h z`HQtixlV4=Sbn^HzFozzU*J{kYrnv&+1Gx7SGTYI0{6DB{Q|FHU;71K)4ui#?DtK- zz-!yreu3Aqul)kAYhU{XUeCVv3%tI4?HAbRkze2q>}$Wk8`{@?fj6?R{Q~>E^9$V9 zzV-|3^Ug2uruMa8;C}YCU*OH`Yrnvo+t+@9x3I7M0&i(w`vvZAU;70fU|;(M_WPP& z;2rI2zreme{Q~>XSo{L}v$J1dpPzn#ea-p>-p#)D3%t91?HAbhE5E@0+1oF$&r`p^ zzIOZq54ErT0`F;G`vvxO?H72seeD<6*R@|@zkm4!9%Wzq1@<-O7uY{9_yrzoU;732 zbF=*d=Y8I5?XBI9{_lU!Z;rgK-lh{AcGC;4(QIYkrxTpl!kaqV6rbvxRXEP$Y}4YV z`Cqs@zIE6-2Om3cekN+>$WQr%Q$BtEdGqs8vrl62vp6&7H!TB8UHlwH++)RS= zvG=Cp^mXG+ag2lFeEoP+=b2JZRWmr{v(fCascB}B19s{hMdI{xI&(dV)6eD1oUdtb z>U>GY=wV?8H;agq@$+*_bAOD{tip~lQaR9{;;4(CXPWzCY?YnwAKp}5nAggVdF^MH z=KipoP2@)ooFN<6CgZ2Z7CUu5DE-Y|>2D6fH9T}{a|(`F)OnV~qR!TYQ)fEDsWTnP zLD|76JLb|{A_us+1;>0*{lTgJsLMxU+*wC|edSecy|E z+)~NU(!vjWtnvew{VXH=Fm}ohT=vsZ_~Cx%XG7-tz`0!$e=#{j~2`yv4eA2EUwB)v8(yk=5%!TZCWuaw_ z=Q^w+w5**~Qw*yWZQMJEcb>9#B1i3~6&!u6E_`A)ov@?sy@eg)H-oTa41CVy`w9Kc zB<$#~mftMGj(V>l`eWamP1vz6*A#Z#qvjBH_*D7%=QeM12|H?r-J~tj=6@UI>xI8= z`NBWnd0WW*H0;$cZ;Lkbt?$v?hdI+Ru663YmgU&d?7?f7oF{Ac+VT%MuiMQ1_x-AM ziY@$l1yjHEi?7%(W74N+!5auIF>P41;Ejql&U63V?QLVhHJTaw-iY9+CpC(-A5iSt zH;-A%_Xhl}`3t7Ljb_}s_Kc}Iqm%VoJL_>FF`n2dpY*rN3m#$XSQwugw~43?J~tIy z!@IH8PjHO?GL^chvli*ESH*5Q!BLkkf+Kh36HfVLZ|p9766Y#{W3TEdIOfY{qGoWL z3yytL#mW6gag2lFSW7F4IH{+q8JzOjXr2-FJwlNKcB_j1h;tpm5$EcH8&vvJXUs81 zYY02IEk&G+-&(?sG1^MlF-9r}`coWrSx@xG*eW~rO;s1>wX$PgZz%f1ZflVrIjkn^ z7(X?(*bOWDnYtHjWBqlsu_UK$OaDXS?KgdyyAos?KF&Y!pZT z{e@5NDYJjj5?j1a!KrhhxZmt5_HbfR=T}kVE5$v7{OB*%lsXMab=pDX#5&wvj6d9t z!j9N>5*&6rSGZjS$8**|!7(no3XW&6-2_K{cQ0J56LAhQ&gUt3kBWA1(hd(d2?&oe6AnS!HsX9+ObB^Gs z#kqxJoy2)~o^f#xJHMhGm$VlMEw#C@Xu%f=Ew#Eh`Mso~y)%{n#Rhs?Fqv%x5o{~+#WORxv9FpWn7HyO~YM6Y;PEjvAu3M#`c=w z7~89cV{ES&jZsxWSi9Yxw|y_Ab^6l3YOQ6SziocvzI*IT6W6|z zv?mBHWB+c^f=?D&Vt6n4J*}dBKWWbtTGqxDwx&Ec@Hs-u+WesS1)ndpqrEiNJ}g@B zg+wEe?Maw*mpG(@9Htru?4UzYx#>IR3 zBuN|Z>73ih*i4qR-V5C1#>LMiQ&hB5Chb&0%Uo%jv{MT$F-()R(^jT87|xiqGYKtoW#*)vMQB;KvnH+AJC3oVhW}&ju5>@y;0=Gb^`6$t-?BDa@xwgw z_vmwW!BOiuD%_kEZmzplX;HwJ?f6aZj8&V4f(+TZM`-|`dijI2|McDS#azZ z$|w6n7h%U<*Hv(h=16SM796ptb0YEih4@}B9A_KGwGZcK8|J!D>+T{J@>4$HluyQ} zhww=(ey(ATWt6{VjWZ5o4kKEqrk=tl+_HjWPV^ESpNSO5IMg_mJM;bzksq~I^@LMC zso8QOheoN(@`A(X3W6iQslImyt=9@UtSIc*x2FAOLObGIN!Y=yEI7t}6~T=t`K>BA z#&$Kq(VxnXx~wkjm=nq;dx)wF>p}I$dgv|s!)^`1VYjB>$ZsvdQO~snN6plnz|Q~9 zZJra%mnOkcmvuyc#Hr4U#OJ!gj(u0fiBH8*&-H{)?A8}tqs+5Df+NoPMDF;Ux3V4< zuGlRiIC4;SaLR5($zcPL1Kfs!V?C(;;8cIqrGvv;H;{cI2?R;8WW%%%wexPrm2ow?(*N#>LO0!$tqpX@u|tH!}GdCHzpE;iWb)9=OrT&lur{`j0JK z^aHn7^0T+_!+KJF;If~6gdfIE`GL!R_7#5E+lGkx!0ngv+F$r#P8?9USRc3plb?fx zmN{~8(jFqT#CK@Y9wxNJaCp)lA+*ekBa`+hp=IokPTFIHmfVg_+T((bh+U0 zd4=G}&(G(}<3J8q3OhbyU;E93cEov=u!Fl=aE$vkf*VotyH;?F?RA2qKb0SKxn9^Y zCzMb2K~)#lgX)jI*X6G+$#Db zPIX2hK5r9t+yhjc_*5JVKbz2d-C&=QrVp z^{f2AWk0_QKa8F71DE}%pTDq|`ZL9Q)BDg5w_bsNi@We@t+U$K!>I z^N={7Fs`Fz2!67neJW|67FufaOwod$6>_sv3WmfP5SR#@z3IM-`u~fBXWsa|NI@b2MR4U^UvH- zd$7%_C*zVMmw19wsKbB^#sZRReuiSfYs=jRyDdBP9%|6Ifa=bwF}p9_Q^?xV^NT=sLJ z@Wa?CKXBQP`i#xGoL$rh&h27emxy?n6JLn>z+IN=!}l%Xm@r4KNZKn)pHWME{`oS- zceT(GgMZeJ+G~ZDdEuY0qxO2CW$gX)Z`9r>wB+WWqoej_p=B;)zqbf2YdqKCR-t9> z_-E@F!|g>I*C6qJY1;fO4)>LDe*Xl2U9`&2Hx)nMCO_W^EqQ)lwBR3#R`vO#@I$OW z6+hsgg`W|gy0u>t=l61&I!BQC_G_|p8~Xb#ab0m}ui0m}3g;B#`g`)}9R2+vIO6|aUjU2|D+wul$XWxR!e{8|z zwO7IXzVbc=k0^fkHtcg7{SW6k(K;^npx$NwjB^lf5|IxvPbxU1ZBTGzlWXk5GxEmYAioV1Gw zEqN}Qw2KKXF)W_6OH{N=CT$0yrDjVd?b1TaTv;Y*I|?o1-zjN33oYxmOVWD1C$@*} zSlPq63Oj1uO>oqs?%>la(bChlPy8rQ=HiyCfFv}zCADEWzdSYPA1nIG!Cani;;Y((0_HYr;AjB{tN z%DUN9*pWv+!4cbL6>f9EF}7O>j*FB@w&o-=^t%_Y+mss;U(cg~7`MC1T zw1Z(ESMUJCJ_bCK^*7ut;O!0jSn~|Fonaqq@V18iUP8^bG3@s&V%yrV$F2Tf{hr2o ztn@v!U~GpJOkaZwCiXoH$J_=Pj=Ak_IOevS;clV7T@Are(+HkiJ+bF{^wvmQoY$FWE*oGU9u?;iqv9VqqHT)lIjc;4_ z(iUWAK7Hk1wbnAv_clLq-+kiciEH;s+NXt#;G@S8%*T0O96!S58U+ME1rU$4F54l+NPqh1FKj(QzZ;SLoX`5z`Y z@;zK|H{nQ-zk=o>sKr(}k8=pHZ~nGmBQ;)6X)__aL ze%tzfveOOzZtFdbc;|b5LXLWNdVNA}oWEu5=F&g&&)*yA|5%®C(0e~jlX!Vd0M z!4b=Cg2V3i3U`O##(8+w?iAdx7OLqk!Hq23-Gbv8=pMlh#ZKH?`WNnAVK=UDx4-RQ zw65nKg}YDK?NGS;1;^ZbKyc%V-9w^3#^FI>$2c7MxjSeb=f6IM9d&tFaE!ydZ%*hF zZtu$cd_?$UJwICE9;U_~$)?{@zV?uH=6BUgG>2hn9VIgRs}kdCh2W)9{A9HpN`K4os`l!aH^-~lt6$z6uWqk?d2`&`Uj6dscny2? z%bVjh?bR=Dj@PnRzq~nK+g|>0eY3_L z`(EgmH^)Qm)h}<3_q12PygBxL%r9?_huf=P-W>a$OM!~FQ=?W3Rl-~XPw8ixGx z=4XsMH>JOn#(00+Z1>-Ie^D^?|D<5*|7pS0^YenK=Vt|z|Ci19zqq`6v|kC1IzRfk zD_h$!re6y?=2Pt)@NX)1-wKX7{auCozQX-bxJf;g((w7CaUFQmUVk>+(O%(y5`M|? z7r`+ezgD>4D%|gdo2tab{^;-@<2EkugZl>jxAB_<{6E9YAAkEV!;Br6cpd*~JafYF zABOYTlY7%I6X)6#F4h~nwm-J6_avtEdc#jD{1WqI6>jnhH$~#uGd+VT6BqN(9DB90 zTg%@pE!)L0$=vD~*=y#UqR%!~w>%ci(Wz1nu7qn_;ifKJ94G92t#xomrA=$vxGts< zeyLr%3a8HBWsWM2v6@c!WUs&UBM-9m9tEe)|0K@og->Ep=ajNnzwqsZ{o%9LpF01O z{uD>mR$btMm6CC*|pTs$%h!dYP39jMcS({mKtoa4RyvL_H z%Mz|b#cnCVQI~}UM}EpDobt)D+G4^dadr{hsFt!cEg?9@VHQy{xLF0qbCQaaXH3O0 z4vJ&1=qTc(o~mYW%4ehGKAXq^yRM=?;#^*E#MxbN+;`L&evDC1VFx$6h;vlQp_i~@ zjOGw_jFHNL{uD=DRuug)w#sf~iBr{ud9Cc2*Q<*Du$xolM-JVD9pk6Q7CUtYBjYz$ z>2H#flfS^tZCrfDm`B9M9GO>eJaf+{`hc6i^Z{-NEqyFdwBYtd+a|=$m@X*%z%5k# zfEN~isKFw_58R@~4|p-*hqH0A+8O$uEzP=0CyAl!ga6s>5=^O6k6ucvPBE- zRkUjUFDLvE&+^3&cm?5yy=q0ljq-9@TS;*2#VZRAyWPb;Moy~;JNCF$1;-jv{W1QE zBlp#cPc_D?8`q&F)#iVDi~gz88p03Un#s>v!q2GE&t9cAF&?~jAL_qu;i4b7 z^^%|Ug&+1BecdFP|sQsaWqmLbgPwf0$w|s9LRqDQ@uw(p=6n2b(&zbC# z{*D%Q^jFL8SYb!KcM|=vZ~8fKITqIC&ccrSyPpY{?eMAc=>h6L@d;0 zpx~(Wu7bmkGuPr=qGr1ZJL?G1?7IgGj=l8|!7(QeEnIwtAU~tvf{zwjVmK!GJ+`7fE@_V!T55Sh(SlDDT55Yz(SlDFT55et(SlDc zT7D+r^(61pjO%O|?HPu<7^a`oi!Ec*-8kxUmT^6TuQQ95b>W}m(Vm_B`1cIR`&{F? z1?@RW>)(&y`xvj-TIdO@^iQFGs-)*=4a~pIK_D2?n!>`6@IAy zeT9pD;OI!}<=t3&P1rGh zZG|0U;B#hTaZV%b=&zPvJ7Gt?Ul;wcZ%!}lSeI`IJMOPD3OjtN{J7W7EbOS6vTJy$ zti35>8PP&Dy(KtedAq{BQ{mpNaPJi^?j7{^zH$Dn3I0H6nOk#PE#h*!tTh*4X_a`XfJOM-Kdb-*~U5 zp8mIhb3L(B=kgJ!I-8DI6h{t!iX50{$9_CvEO7h0JRwJa>I^n==vWy)<&!x75^>`5 ze}Zdxc-H{Nf$rN&KK@wu*ull;`VS;V5w@Zz(;WTHRfoLq4DoI-F7 zA2w@K3XV7ji*dlGI&TYZY{hOb!7(SM60wk<@(HJW@)@M9@JXBpidaUql%?q)!Hp=~ z)WRp+G=k&v$+UuFY!$~i)Hszp&oldq9H_ObC!F$0&Dx0^8l^7N2@apr3y%Db7IBip z48o4R;4rbqh;v3^2RD=8826b4$JovyIL3BX!O@?}kGjkz?3feEC!al4U04sQKi0$S zqCf2B5FB=M3Xc5d5*+oMTX57&%?a$z6LW(3a{OhLgW~Y1YK`53qCfIecI3d{ zk59#VJNdz8>J_Wm3dwZ{l8x%u;0)E+0a%!Tatc%fyD=Q^Aq zw5%O}MvF11^GVtNI+=E2`+H}>HC%43i{O}>AC=hRUd0%xb2#8QuOZqIr{YGmgwoVi z#KL$eJO0mWH(>|Yy>M~eqG2X?C%7xP~=`RQ54#b5H@>hDtFdQ>>4m@{;v;j?%08M!qQ*WJm|luy>kn#s;h=x;5-v9{M19DCY2f@2R`S8&+1{n>l}uP%<* zch(bjb>DTM9osKg=TP1Lt;eytWd4m=jwU zF4hNbo8-sqo}Xu!Kieklc0x;h+b3;*p(Ta^NxOs4GB0*a+MR@!vEMmqcM)208<@1a z3N3RX``t}wS>w45y9+IAXHbe^kD`tDY~t-w=0@bG{eps{kHNwxb_)wT>OMr+F@B2* zJI280Oy2+KZ*gHqf3^IU6m}!M8*4*Ff9#t}2|L#1p2CiM;xfVxpDMr6Eul1Z5_Z%~ z*)c}LL@d;0xZtSu2*F{uoQP#)iF2f|qb{QaNA9XW@>3jhceL;cH%4&OWo(7pOK{}A zj>sKuUBOY$y@gM>eFVq;pzPq39ko_=aLSH#w1&t5ZcV{4M*E5!;Pw+7W2F8LF8%E< z?C^Pj;24Jk1xKw95*)D{EI4Yd>_(P3ui|8$9U^?f9V$3-I81QlplXKQR$`2Z^Kjvl zIdO#GsI}^kG1^h|2X|z}=TU;g?&u14jNsTqjujmH?s0+}UB>Tt!7(RJC|rErASr*YIbwd-XgS&;jKw~o6wH-!dbgLY3~qP*6p21>-Emhj%e=^9DDTLf@81s zGc2-C_ULzBn~>vPqBzFH&)dj88QaeVN1XRm;=EUI#Cf0K$nQvz-$)P7+Wo?gd&vWW zV{WMa$WL*M(SyY&Yu#gjd&s!BmL3-UQ>RCSAGk-8pT~qBYI8=ZO^gTb@#N1NTDm z<8}A76Sd5d@k#rl&=TKEN&B+U62mJ=`>N0~FJ4R9*M*j`erTINFb z`;O4E#&aFs6D^fSNc6|P>F1r~SXh@I3p?(we)dVW!>7t`w9BpeIVjnV znkhSS_(a6=pu4j6X@&btaK!n!;Q0NMF9b(jz7!mGUkR>%OPEayU*KP~{ysa;cwY-U z^7}?`{Os#n!D08E;E3gW!7&cooikyaJCq!L5O(DLqu}uQli(QJp9RMr^h<^NwZi>Y z;eM}he-uvb^?w={f426?hZEL`x%F42kN;J;zbo876|U*GiO15W!c9`Rm@hf_c@^>h zzDz3o@H4^5D%@Y6Oc)z>wKFIt7e1L6Q&hMqE8J9ti}~Wy=S&>mw!#nVPQ}e0K6SF| zS=NrfaP9>6q4@LA%=M_)jSHNg7vcNL6yd!+%h%IbEN#R4`RoeE?GoQenEk}x&Sy_} z?Wo_hCGWL@{VL__nVPx(&X1V2;MEI1OTlXt+^*m?3vLvQ@7W7pyYTZD%=r-W6}(>I zXDE36f~PCEPr>sRyg|Y96ue=@6vVZXdN-ojq}^5)pj;_%CxAK!s~BIBK$B;U;m{r7<20 zH*^2RyO7W_){7J^@hmE|aEli$&VxTdKkx^_2(F`gxW( zmlk$#%T&0I6|Pg^V*T;i*|?4t4Y-TYQi~sdZqb6(860q3MLeweZWXS3h3k{-HR3kGd?I;*4Cc!o_i7-|_P~e4N0O{^FLczqdjk zZH;p+csbE0b9ng*w?c(mQE=3CCBZSjR<3ZXRJc_O7i)u0Ul(x>uO_tA!q3l$wV9^W zG42b@tIvLzu*U4WtCv2(y+xn+nf{HIKEX2>&SQt4HIm;oi{BU@b6cIkz@Do(`d>@< z^3v5 zgAY$>n+remzeR=HQuKk{R>sBg*jo7EzOjw)!@AzqxEO1{GLPdJpxrKMx2R~hPulG& z+Kykhj$we%vS;;Aes>aD*6VXE0X>!Oi)AxRBXXkR0JuzySv5i?(XjH?rsJ1&T!TvbB6C(_x=I* z{d~^Z%zkEOO{}%vcki?Jix>mDp~hnV!-OB^#9o52rw$ja;q9_xgkZe(SXG=0@u}{q z4_mil*IO`ZxQ1Zlr+mVc&mEi7Yg$|QB+kBqjVN}#1miq3Qq&AKO0ZGIr;3x;sfy7D z#n`*n7jaTgRWq3K*(h<27CB(oPmD*Ln+ryq8wtjqukIyIk2Vo@u)Rf`BT5dN3OjnV zkFcXhDhI|>jJj+o#-q2&j(t?sg?X*)nAZcvc-ZYL@*{`-!jAr_-eRZj1jwGT6E)UQTABn2bTSuCj8JljpsUCE_7KtSELxOEV^;767Th8ZbU}yFA$6|t`a`6^Lw4= z*I(5AYGFtJE*5s9N)LR_WS@+8sjy?bT7H)cJL-Ln7>|APN@2&kyjIw8j=Ead;Zxt?@add{7iDLdFj;(ae-sj=%s4*0xYu!e_c#|?r}ms=`1sC$$%9`~}2^AqOxvd(p3 zUib*Yr>I7CUus z?G7G-9k+||==&XlQL{TM?5+yCyTb0NuzL%O*FwZ{pRxFU+~`j|vzC8I+N%*90L&Thdse7(6-gY8iYO#}G zTiPsZ$icy!nDm7F0StN&7MJ(i?#vA0}*>T(( z{=4P=*zl+9cun}Ep0kMjxDGf|^oUs0y)dcS>ta0EX~HMDzfrMManc9H@TqEz-J4=O z@>6!?aHzUejQbC3{4I;c_e12jMQqHiE5-PHUTxC1z2hzC0XeDrClcGO z!Y4IUYX43pzXQeE!R}pQM-J}^*6?!J@xEa6Z1GCZ)cpk+Zy8}noQhGi z4@4}~tY-I#=n-~`5zB{SJm%6dqSn;qBVh;oxWYaWjBAjtqGtH~RM^2jtFX@nW4tda z?90O9dc^K4WAQz`uZ16C`=-Lat+4L|WBeYXPRyV0lU?`bDL($Uy0;ed^@n8V2F#xy z6YEs**}1|TG4Fm#KAjQE&w@2v!H!=9qt@$+y5LjYuM4(u#cmV9sLT3-k)QGjQ$BgE z*jM-@&MgHS(Oj0M{(>tw~~gJQhK+)TttJyp$M%4eg*`Mbyg zyRF1{#JPiD#JP=NT>Gf|6w#wW!VWf8#7Y0Q6L$1yoUo%uDhI|>jJoV3#-q2&j%yE9 z7v{CHV_pvy<6-xQ$d4Sh7IySc^%gsIpDD)svohXaf-%qj7K~W-5^)lXx+kGj;@ zBo&{%MV#cP?#Dwc>i$2hhe^eF>;-KF8|CG&qn%*bohWL}x}8kek;CMI@m`Y}kNgy4 zPE1jJ#_K0odt>qXX-Y9Zb(%`}flZzKOe6eIn+r;9Vmz>Elb`8?AL>7SVbKq4hU8~P z;fFm|`GI9WGYLQRPWgdlKQjwIobS#N^?|uv?AI*94|8JH!eV`3vn4-XcVGWempL+f z(w#%-65pIjcP^nz409*ld4w+WV&0@XpU|cEjifuj&?UDHNp}IE%UsBQ7Zkdz@mz<6 zgf44m;S|FnMK`|xNW6E6eTCXzB^YBYDtuyhjj*HciwQgW=l4{~=L~w_b0#y!yFu77 zUM;_ygdO!>T#U!Qd5f@PT`nQ)I7i(k?C`1b8|f+A;rC<8xl=P`2Yc_0iS3A`#+DR0 z;BzU#8a|92OAAI_?yuya?tjF1+$Sf_U6|h|C)b5}?e{9mjNFw^nDWV-e@yr!7QaVM z9*;T5{c>UsT<LlYK=o`k)x=NZn_Mda9bil+T8j!;a-d4%n&tn-J$) zqMpS0ieOxKs{6jsqt}ETtfPpN{`vica&PI;^1_ZDsT>$jG3xS;7?0j6JN7J97v{CH zV_y3`g>rt_brSiJ1NU}{^P2vt-eRZj+e7Zly*qI%WBe&QI*VAwj3-S^X9+ueb`f^O z(p9jBiqCF>F$ZV+(3M)gzF{x#F6_n>pDPH)T>4t{0lWVRJLdd~f(k395uYc^Lw4B@6^%i!p^(w4SVKEQvyuUGz z4NUetIMx?_Su-07#(vm0G0%@2@BY%gx6G0C)UVjZSXoc~$Fm9F!QQyYto@COZFhpV z#}4J+!=H=yG{k zhs!=sm@6(1GTDFpaCuvM^dCQ5j&sswKQD1^a@o)Oocmq&b2jH~m;F4=x!PqvKXHz7 z`TySorG5_-*6jdWTX{dA51SS?iK~=@y<+q6Y{GqBHWPm6&lZ9)#+G6X=FwJV3}W57 z=(Z&odqj6&(Z%mJMP^P6D)Q)JyX|*^-K=*7(yc1wlbT3E~#66Za|PVFoA8uMD9d!HDC z@$av&2P*7A!I%p>+1$)EWW0xj9qi!>dqgnCd$htHD=gLtyT^^iIrN0^Lu^l0*i#ku zv|x-sRMd%U=v9B4xQ6soF|L!J5#x<^S9UbM_TMefi!gN$31WFx_#_r}uM4g}w-Y(w zvo@Z(hXms(M()pvILU8;O78031&pWel|X*`RPs|kiSv08Cq7>gtl@35<3+)kpXwe2 z_*C~pfSoMnAk4ovpXUp8@p~6!Mt;gCO!;Jg_4^oPpTu&OU|i4h{rPz9!906O)C~5r zV52+)J5-#!FQXWJP>gey-#;MdM?F=|V9IC1UD@%9$N@X`eP`nI@6+d4h||A6pBd*) z^*wNUbtFchh^t?+`8rY zh*;EjaLMnDN`8#j&RyFvsQmc9-?AKhY{qvv-oyDFll^r7uLa&T+2-t}&i8HGN%S-xL<} zz|Q*{^Z3qW&x7Mz;g=eGUs$~6h5b2 zoM3;JoRI$#y43COqKo`b(T#N^p2bD1jG-9w;a}l%baUC7?k(qpI0uQP>5t}lh1|yE zytb(0Bo#Jkg|)4)m)>iAyq7EN6~QPTBGRd(5nBg*jqwty`H}iZQpQN^!bR=GHvlxYTklhN_#5tQ_ybhXHFy`3|qGqre1smmM zv_r+oYaPYtgJQfsoLR(4Jyp$M%4cJIlub=Di5#$-U5rPZ4Z(CxQ6 z4mPuhlm5*s?C8-f!j2xP92ieA>e4}sM{kuK=UY`5=C!h8UN0=h!){iQA34k+?C78B zEp~TR_9gZCa_XWOan2^<99`zbQkA_`eZG~~k+X=;M-q$r94`5(|38!YQX6ln_dVd2 zd*nkUhuKAb#PZk$6Wfg`c5?_j@|#nzA;oSk!I*<{3&!6z<`ImwIj>;6CZA6*a&HvY zzQneiJu=>6#yXlzeHS&^uWyhSF}ah;#J;e}orAoP$z4n)js;Eb8sr5`_G>A6+ri}S zCR6+QP4=Em;}3RtZ0L3|<}zbAipQl=p2lKqYnvQnTg&7a+nOfF*w!#P#zEvK>t%9l zb>GTZjBQJkV{BWP9An$u}YcA!wx3LK5TDt?89~@$3ARpa%*#Bn6Vh!P?KY9 zdzu_$8)9;dZ4Z-UY`dErW82N-7~5czV{E&c?6Il3A80Jr{eU7<_x+1Z-S;apb>G+I znA<)k$K3WdIp#Lngp}*yI@7K_<7hHr5!|a%~)4(Oo_1u5YsYjcb0iKV`=;#ZP1d zn^DuTf^ki8Tw(G3TV8vdIG#=Tobd5QW=@@U|bKLTVdxF7GuNig7Iv^ zHQ4#W5B(TZVHZ}|MS_vj#e$7&E_c%XHp=ka+@vWx8A#C)IUP(JC= zeaX&k81H_;sLKO_(c446cP}mXbJp^AKTXW&!-K+S!`1D0sKOquutzHF(F%L4!X7Vd zissP%`+uWdeAZas@;q|CAU|Wge~_OxnY#JYPnk?_kg1)^Pa5Z1#pNeV&b7whbAqvs z4j1QK7@vD@fA%F^>)TE1}D|^J1Hyu8Vw$ zSRb&jMLevXZz}BD3j40Yz88#GeyFe?3yZbkJnElq_k5B0EPK>t4FCLitj$kiOxEMi z74}Po{VEvy=Wl}1)88v>Y=w<0Ear<(pNnzc{2_Fy&6rdh^_h6~{Xa!KT!Z{282PKu zs8etCS$g)3zlBfOKO!gCzsBNz)AXlJlyn`y_-!g|k_wx&!lwG5_3@^zuxSLNhHb@o zuyz$TnPAk(Kil3~T_zWH)NG0hYhPHLm%L_}(wNUnrD!)@F zzta@I@w$O)74;ct&IO7w{NA*wT-}brmCs74&x%ri^?yE* zALFSyO_%C4eagxGGsX;w`F_EkF=Jvr51EVpdC6D{?=}6NDcQLub(*=b_WqYixnB$MT^WHx=4}9Wnq(JZVQi4NL@#4;PV_7?b7JixGbh$6 zGIL^0lVfgcm>hFk-Q<|tY9`0rdYBw@Th-*4+bSl<+*UR@=C+c_t<}A+u^8KiCdb$| zFgeDyzR5ARJ|@T5)-yTA*4yM5+qx#l*w!)GV~cxx!=JLFUxoFru#F0fdnUc|&kFl_ z2$|0YN8N@GZijzP+jTiME`HU%v0wQ+3wzeSMP{6Rip(6@yU5Iu(I&?}Y-MA`K5S`n z?86o&$3ARsa_qxqCdWQ(YI5wu0Fz@MHZi%iIkKa%7~2je$Jn+vImWh~$uYKVO^&e* zGC9UJ(Bv4~HYUf|wl=x7wYitExHg9snYs@xGIigx$kcs^$uYM*OpdwjZgR|RHT~Q|!yZul#Qgwvpcn&okg@n0;oyp&Ln`dh z3Oh_N=Goze#ah9RFc$ZmBZVLS)^n6#_&K^V#xZG(V}&1TbzFrVFUG*`1Y@zrCkj8r zdXn%1JK0#Q$0qR|RqA`Ljpw~VKCg@)b$Jc-v5n!n$e)xk zqb`2VPkt{Dd9f#}&jZsp#Tb8#@HyI3wPWG${ddcr!qn#x2f4Z(vsXUnr#?4HPZzU2 zFONrljHl{!VXD(bBBwF#+K!6_qZUV>>l!WdC8wXxpO_K%CBi4K+b$JssJpUb@iG3p z<#-SIQ+8Y?>^L9KE8Y_tGv3PXq4u)>qoq$?vo0p=V3&*W$nOflu)DItt`dy8TrC)V zxJIy%CBJJ0W1e4ESezrodA%`TcgQzXbZ<<$Hwj&8b92!}zD4LV=B>%^Z57?ylkOct zms;Lgbdm28y43dWqKkZw(52S*7G31~iY|Xs@p_W?{l@&d3*84z_UkUjd7#+R8^6Y) zE)N^)8hkxebUBxF54w*eKYlGo-j5mU9CRN|x_-?&*na%zZfn7>&dgr;_f|LYJC7lXRaIy2S8Y(tW<7`$E!vQRq^$my+(wLYE%Cl5}4cy7d3Gr2D$i zW!=7!biLlK?O{DDd)S-8j#|GZ7`1)7!rl>#+Py0nwR%r5YV*Ee)Z&A}SSN8FerPQ2 zVINg=KTf)z2wiIPY0*XgOz2Xp&y(LTD!N}L-LHf$wfws1B7Y-vsqME#7x_D(ORc{z zy2w8iUA2e(Xe{nwKbsu)u%C)8y@`9+ug1Dsv#8-OMOW=%za>9$4;yPN?qR5a?EWylVfhvnjCYR#^jjW z)F#K=rZPF^Hl@j)o2vVK#$s&qnjB-B$K)8>+$P7^<}x|PHmAukwmD3WvCVFBjBPfP zJvP>>OAY^J$7X#C`XS^1Y(y7v1@j?%6_@wbh~MBA+L8 z@w-6sJEo$$VA8!<=(08rwKe6rAzvnRS(^(LzsOe#UDoQtMHl&+qRW2h)h4%LmAz>Z z^W%MqcGPQ8!SK0Qg)J@^`7a?D`7S9K`7I?FYk29x)O=aSSe#qSR&Ng)X&TspulFTy)hreHCNzoZiD^ z--8%u)nZF;;yHbFV_i*`8m?A!nTx*HP{TEnAK%x=do5$}oW5q#jpuYe_mAJTldie_ zZ_J*?;_H%L72S1`?z%#kJbNeI^@J|*^+~$xS9CW>x*G~zYSuUD_7l4Fuz%9sNa!+G zHcq;m2wm3gfTZj7&gY1RKV`?7>3pL8UyBu*Rhu>~ewau8G-GZi*wA9PxnQFU+d?p| z8MYLR@wO6-Ym2P~!)_bF=*K|8h-FY=GnKggp-1`mjN{M2O--v8d6putQRJD6yjGE? zDe~GyZWOs^k!LG%uOiP^3u@9SUP_ z^rr_Ew$pev;Tm{H;fMF!cNUECcM*)7b`^}A1`EdCv72Db>D>im-tHmTwk6*og6&q= zo`R9zP{FVpCK&e+*h?_t94;7j86g<4j1-Kzj1mmH(Si}@-h$DGeFVejzJfJMJ@*rg zJ#c@)$l-trJ5Vt8q=N)w4jwERIUFJwc83Z^{|*z3{0uRl1{_evTC9+~X? zY0mXW3C3JHIQ)r54AU?B{;<8va6K@iqLIB4c|& zks0g!B9qH`Cdb^)H96*Xj>$2%vrUe;&?M#znZfBSrb35JSnA>S4dv2=kR~d`3 zU1@TR?Fy4)Y?qrHW4p}c7~7>L$Jj10ImULe$uYKzOl~ccB9ELwi`^2 zv0ZO+jO{v;V{F%&9Amr2@-{fMqTwC!=XNu9p4u}?0^0o?1u{bQ84DV zfB!Jgdwl*R>|j4v*e`-H-mew*TVb)L*!^xSzNa--_+eg;tFS*R?9U4OtHS;+EXGYN z{}_w=Z})E}t~alX{;iDB^!LPejM1jTCaJJV1tXTW71pk>m@jtz-8;`0c{1ULo=#q2 z3;r;1zSz~iGdG3sNuAnP*pwAERbeq-eEOV;{heCqHarzO)_lDAI-oILYL@fbWA9Zt zvi7DCvBIXUu;~P&rvAOa91E=N2d%Sq6?VG#KYh41oLD&wi%&ulFAeijiMerB!snJxL5UFZ(?7T7UI(M6uK=*DvbW6mY~z~(M~kmpH$ zM%h^YX}Y zw4h+vEmUC(3&wu7h+w=9SyZr*&FM8QCK$C}ys$XOh;s>JKF5%ktmrP4be9&o)MlBY zi@dDRWz6N0-;NdC<&$nFp-U}07hU8oLYLZhExO3vgf6x2UUZRHD7xzJfd4br(dvut zN+$btB;%}DZ0U_(Gg6mTjCD1c8m?S)S&M$Xi0-P%k6)XT_iD!c+7#U$N!PF8S#L|( zTJY;qbXQNhhW=w~7>n0TYgTmEO1f(cUGnUibbAS1;#()_u3ORVopjd|y40*s(p_KZ z(!&js?uJ5_{`XC~{e&*-wtv#~dgpT}by%^ohixS6sP)Ez;d7G;8z309+f*=WwV7bl zW^=))#TJFJPU1Y=(pcQXwyNlEopiSmy3}T1(M294bg9+0$?tX*-R+a^4nmh&?pSn@ zcM`hPcITptyo=DK*1Hy6|5;G`;&Hfofz*hW8PO@6C7%Ce2?Q0lf4JLCw8#O-UH-=Opfny9B8ul zdYJwAae&G3J&yfNj_+~oXR^nw-dj4^SbT5kq#~2oiA83t6N*eO$D15;JI>^o+p#9c z+>S9h=61BnF}I^kj=3Faa?I@rlVfg&o9wx%x}Re##&)*JF}AZzj0k`4J*QFJ9nPB99xnSgbg<#}&rC_Y#s|t(rkT|b4 z7U$MA72Ru-?sY319RJqX=E-fbgAKeMVGngdkr;wAo=lqjl3T+ z=KC7D4<_AsP9K@h=?^DebMv2{M~ubSC688gA4|HA3tjSjBI!OUbjjtZr2BM5_nD;o ztk9)q&n4aGg)TjOA?dy-bm{*~N%v);%es9f>3Y5MIilfD+3|5YpQ!&+aA>n?)2qb~ z^T?luy;fnb3&wkPZwQ9ZH!JL|3VT~H#(Ss2-mS3r1mjxfeZd&-g9`hwFxIfw6T6R$ z#WnJ=@I#$GsjyEg>@&d_|MLp_qQbr`Oy%Ugjyb9ON-&4N5`JJ`8}oAy@;5~{&P($6 zw!*#>V_@g`$5_?<8PL7&1MghRb3N?-3j40|zVr7YHuC+!Sj_*&q8sx^_os@VpOc?o zgzj)p&5mD-F7j_hSB?3*@IyRfiy!21!q1TAA)EdXj5&Y(XRg$;=CF|;Pt3;nQ+E6* zeDayazXTi6Y}WL*V518AM=eHT+lDkweoz?!0AwoMY5@oC6i3)@_Qo0LQZYVtYAgJ}+9f}e2|v`vzek?WOR&k4pDBbN>fgSwm=|oy4=nqc zTKJ)N$`362nMU~G8pXfApX&p2yV$R3g&*d`bcMzGz@|@ryzY5S=Ew|5cSfO0d^088 znT0Mf%#w6x6}rre*^=(;LYLmpk#y%2y5u%j(w$rAG8eMnd4w)&JlA1fq08EtFU8O( zx^eFyUjJ@<-XEyF`kp*v%rAUm=l3YcKB;>LVMqV`eg)Z%9{8NejPd;51(`8kEkD1H zLAImb3yAU9H~pRl*^YI&ps?dQ-S2OZ?eMAcqP{^ zj{671xkSwt6?W8RF~P`PjYocpF?SaiK4D7;MqQSyu%!eecmJ+`?gNbP{l~gc&!vS= z*fN5#KPWqxvZL0@4yNo_M}CiloF9z)CB*#b(Xt{3*m8o=BlR6N#_K5T@VUHT^r4er z)Vi}^#L`7DYOU;;^D0i}Sy$l`)=e;S=q?yJsG4EdBtG{=oGS>Q%!&UAMy=I&^r*cU z54K{(=SqTMw{nH8A{cwfs)Dia_7IG{bv4156RQ^%uXl)Z4P)_r%QY*yYbD*ag)a5( zS#*(m30-1XC;45sqT4&^t|xS0^}c0)V;yZb zLU&`6RhF4{UVuv$ya=ZB{C^iSfYpNq+VfeyIO`g+)KG{ga;qgrEQBl<)(~ehw6V=$(oO zmi-(g{4kHZi~7LaF81qS;fFbKNMW%)utSp{ue-0EsLLEVEa@IDbcyeXq7uRbM3~$j+!aEhL_5YvqUWXZR>2osOLF?4Q_VXbgp0n3p-CR;yhn4 z>@E-tyD@?h=Y@sEJqC91csAj0-WLf!!%EzjRM@40vDPmWjDB7&*sf)~D+Hq!R|-Zg zt`dwnaw7TexF<1xn!iKdTUc9^?XhS1_x&5Z% zgufU1Yiq_byHF=@RSZRB^6>y_8~cJL>SVVA#D=u(^aiZ1f|LYG>7koHZ^h>EXXgx9MN6Qp?wY^uLXtHPIXY(iB zxSuU(%&%+c#R5fl3VY6vKFtr~f4^k&sBe<{1^Gqe{e%31$;?lG`gxOi9f3S7=s#zi zYaW04S(AOuw)Y1+JZ^fql*Q@i4fK~Z+3SqFgvp&u=Jn3vCVNeg7c<#w&-ja)?6pT; z#AL4veO%aNuRr-LWU|MmzAv+aF^`q8x)&MSZbfFSu0tu4wZF!Sp zZXHdIxh-dM%xzhdV{Xfs9CKURBS%(dtj-8GZ$<3g8t zvR2VWeoE-#ckSf&<%(|4r2CrCWo@keNpo(E@?T#t^4&l%^4m}_)^Oj#)O_h@EY7X|72S=J z?#4ow+H6vEkp~D}YPD(dyIDnd^Q60l(504J7G316gf6w+y67TrBXp_tz@m#hsOYNK z4BHy>JqX?HP4+#AakeYA^d_ETcQn@3bgAJEMVEQ%dkr<*Df#hzjl6d;7SFLeC*63C z<^EFm-8Jc&+yBNKY%E^??pD#=J?ZWtbjfo_(%n<&65r6IJFKF+SJE9Wbg9{hq&rgR z(!)_nceK!@|9dCheS|LScHgAy_0H#rhCgM;|lph*kKiRc!eENVMhwa zyf{iQ@>BN%gKN?U{X7IePWWxO z13Qivj93>J_c3dXw`yuSLD<1gtgw>`i+Nz@{f&8?VzTGKakB8sIy|+)POGrf1tb4i zKk|`VUdzB{tFW^~-LN~OGTxbmasKkru{+CHtj*aKKj&1~xfOPvV9cTO3yZPAE-)77 z+Zf@8JT9!Ti^Le%U2H7oe@XJ=uc7?MkGIl(J=mV}A3rWlcJ7}s-cEL&A34q|zQ@xr z?#E@xr$5Kv5-t~Pge%x_gO4zZ_Uo9BtXf+;dOEKzlP4TJb z#kIy_&#n{WQ>W{NAJ`4a&yB(lwdr1J6XSv1l>FQ*{80Z}3X6VVw@cZ_m3pqM};oAJ(hGI7rM-a?Dq+w%Noyhcv9%HcAiQx zJY96-bqn$K7yAmeUqdj)ct-fdZY^O)-Jcb9^slF|qX#}`@_Cu@))97$SIe)ru%q73 ziSgJs`v^PM<@3Uh*Et&qJAA7Ac+JyS*ikcOM~_|*u~3&61*6t435MMkA{J`)vaq8r zuLwr&YCQ5&jJf-&@CkcOFzWJpg}osdx$i1+hYc2tdcG-q!rl^${XyBmlpVEJb}(hf zI@(d>0NY70di1u)0rrky^hn*4g7MxJcKCcxF#7PmVAT2p!HDHU!Kk&eW6rBMnP(pf zpRkVwBZp4}BL`J8?DiEsBF;~RPv*pDf>CQV9zF8=n6BrecvXNw2#<~VylN4R9f8zVTZIhq)zVBqlItRb)l5Tw8_cU7z-Gc7qN!QSS@l9bYUURpv z=uVk*rxLoYL)g(jzpr4nqX#}`GGn}_gdOA6^7H!*W}no1H8CFhrr&!o+p#WJ7j~Rq z{XT@*4xcJN&b59|!fZ#)lpSmh5esb13R_Ds@>^Rle4g@Y>v5i1VW$bkcs<2V{*eiDyj6H7`!SK1OU=80w>=-N}2u6O&CrtU|{MkH&HX#?t*b{RdI5Sq8NQpj6HT5 z5hwLjHG?Uijqx+WKmRrHH2`)qi1CPX4#9|XCcy@m@znj#=+P{~4mL!@N&jXOcJydZ zVMmWt4veQ5b(u?yM{kuKYg*NXd9Cc2*Yk<-up27!BZnD<9sN_i#ZKM7j`4<7#@kCU z=Gkz;h-DEGC$Xq|zrmKS*exR%IVd}rvSTid5IMj`3dVd<BQuA{KRDG1kLAVm$VOeFbB^?o-%^_x5A5*d=Lq44 z`X5~SI!^duP8?rYtPkvj z#4zLS&gDMFVRPEER}30>yJ=}Gqtp-b=2OuA7FBW znG4zPxk8sUp6hU)&}HqMpJKS6=*IVSh<7)!uTc9f1!Igc!Y6iH3p?t5p|GQW1BD$u z@Hvy;4`IA*g&pJ7^4nh6QSXbyc+j(GQp_J&`J*KzG{rOU&U^J!I;-01p6=d3R6Cr^P`0war(Wt^892D z?jso2JC}=^!LAUD^`PQpUr~%cD8@QENW@7!RUcr=Cv)OTkpp&ziSdZ@IKhbX2*J4S zRQHpjM@I=e*i|A<`ge@5qeoW@J9?yYU_8aB%L!sUdaLZ%vs7J}^U98SeTo`@>U- z`E&I4>BOSlGl{uf!++WFtYGxvIl;)``NVoQyBq&o#YvA|NOo?(KJlVp*u5keb}tJ? z4zCDC&0ZC(F+PK)rq={x4!$lJKHm_Gn!TA=k201&sG5<(Tgk3huzNc(N6f)@D(u|~ zdrvTOc)zgtI*)V5hvV6V&yRf|bm`egMVDAV7P_!ci!Sx}r0BK@b-?dug0UBWE*R(h zF9hSf{-t2py(!M6?C)O*JI}Ra-!@TzU zgXQ|b+%EQOobbb(_@l5`AK0JCkJmkq$sGAB>HaNriSM7J`>)U?hNdexeo1xE^BAX6vI?SH|}G^J8AixZe-M6 z-G_@YrWQW2`&H~!)O{LZNB{hOXx^_F3q9~TlNsZU6LySO%g^tVmVHw1X+@mun|{x< zY{$BsPS|nY^82S{JAA7Ac%P@9OInT80P{e$2dEf9OHC3#yP*qG0uj`IZo#1Y?XR0Yj!bCmt&mEm>lC= z+TyEn36pc2?C-TWS1`L6r^_+U?k2}LyO|v0>}qn1)8!av7n5V0 zolVYhVqc4Mb+e0cx*X$N&Eyzo50hh@tD4L>9)Zg-&Q(l~ajtA~j+1?U*2?@`&+KBH zF2^`~n;hd@*W?)IIwr?BU5;_~GC9WC)8rf{d)Dlg`8mMsVw^6=I5#ml#<{V{G0u%l zj&ZsiTShhDg#^Q2}Wjm8&EZdsgS}c1Si?O&IV;N#{jAaj#V=TLy9Aj}g z#`%dsZMSdKBdwOGzD7GrTa#&WvJF_zOz zj7_xwTj>FcxERImU9n$uXAmOpdXfYjTXmff^oe&k6^q;m{%}f2h1lJ^=cFrUz-!>{Kn$< zSvyp87f8Ab3SDZmP|-zRSm+YNBFXQf72U;>?&3n1S}sv^k(U&@)OM+&i@dbZrPj+7 zUF2npuKGUfa>hE^%tLp1lb1J{aXJ=Tdeg-ib?I!ZtI5=`Q_*EE#_zLsNq#z+9eH;% z);aj?nsnp$Sy!^P&@Je8Pr8Qwi*E&E@%^FyRdiQOx+@7?@?1ITt|D})(W*(eM@4tF zq`SJ%rDkg+-8F?SJzOj4t}S%wf6t`bOX#w0*Galw@4SbhyRKlo#_TN^uj|$kdmpc* zI*C1#*Ds3EleL8%z1=`C;#{v1XCJ|cbA7?cudm3D_artDcAQf-6pXo{#v?z)=uzL| zleO+^3D(b8Tuc4M_|$15;Rm*H^0SHXLv03@+QfKZ1CpOjg&*p_Sz*x+Z1d!23*m=3 zrToCMpDl$SdZ+xrvY)MlALj9vqCPOUi~ZVK_+d_LQ&_AIY+&-^b@w$Mb(tfBlJ2%b zm-x0zy4wp~V%Q<+?kIGb7ds{0orNyF-zDkpDs;(haMImP=rR|w-`#~SYdqIs524H2 z8Ioeyv*>cZjq7z^v9D13!GbZyP~j81-Gv=>A13VR-wnLG|Pn93%+P#GxHB)xmd4uh^RGgz=msQvU-@A3o*BFd9 zLgWVzE6V;xguVn5pCnA=e% z$J~xIIp%hR$uYOXO^&%8W^&B!P?KYBhnO65JJ@8;P1XGjV==bVO^&ghW^#<}RFh+D zrn;c_1&*T`}xhBWh&M`U0cDBhe zwzEu*v7KphYjfgSV{uMgQ)K4E)kS7bTvcS|#FZw;+^#S==61QsF}KT1j=5cGa?I@# zlVffdn;dhy$mE#Yg(kOF_uGuc*lsmB#&(OzF}9mcjgrf-sBkD zbtZdk{+pvw{&9PS-BDq87N)L6{65t2d&9?-y*BDH{-MTPm*cMDSM3}3+x{Bo$bCg- zj@(;h=EyxoW{%u#a_qxHHdgGzgC@s5JYaI{!~G`5KHO(=?8Ch#$3EO+a_qz1Cbu?6 zo;DU^d&=Y(+mj~8*q$&s#`d_$F}BA{ji&w!F}IgZj=8;La?I^TlVff#m>hF^-sG6ub0){!o;5k<_KeA`)%`+a|}@-ZDAH_NK`mn_3%wFZ;MQKC0;YJ=3G^$z{#C zcaH7KKOVJp;A1NGScR?j!^Cz&i_ecsY|$s|lM4G(FycH*^H$iHM?w5adJ(q7=2KT`QrEf&N)y|RWq3K*(h=TE^@%`9x)zq`n|Ao zEX3*ezRqm0KV^ryhci9$`+;XW*jN!K{quW+XFGZ{PT0{Sl>_4`MqT`#;dwlItL%om zx*e)6%xh)Gy!QKp=kc)nL*z#e_X<1ur+SN>x(7Go{aG3BFTolfx*dNDMl3IhIEh8w zzZ$0Qn+#L;O(q9r2UB*;rGG>Yuzv+(zNqnFYCP(q?qke48eO^PSJNcE8@IfMAkH>| zQ8P6jb*ZsQDn54;agv|9S23}udn2#&f}W$i4SVpyc;#=V1hXRmz5eG0)CV^QIgo=qw2sQY5Vj{Z$8?C62dnS9P* zylI6UI z$N`^A3Dy`NWmD7Af>D=wDmkcoLo?o@6}!a*V_wfM7`ZE-Fy)guzo76*oXZHt99&p1 zu6LFZHG?fH80$gB$-bf(eNc>bw4{iWda9biluzcwav}%pmKEa>XE(u!v!h^KcdC2h z(xXno4%Sh`N&mVCJ9@Odu%kyR2gXy3x~w3^qqoYAJxkSvd9Cc2*DH(hust*~wt z)?F}ay@FuW?7T|N&abd5EA{-J7>`=7SYay(Mh+`i*eZfC-l`SWLomi$t-@BXur&mu z)>~B8!V!Y8lj2&Zt_unn|BAB|TC$VfOd=iVg_b79@gZQo{K5OHt zdwMdSV&vXe#7TaCR&rPOl4Lw}?@IFPUdd1SB+h;!PJH$ktl@35VTUH z?^$6(1!GPO6O3!dy#!1tZRV3gi1+Zr8zo z*)jaNrly8H`cG5S6Ayj!Q~ZClpLpnNlk-|4kA2ga16;Kn%=`E{VD-ZGA8*&xv__Hl zOMd)m*8TyBtzT^X;lRZD1a?qjjyNA3EEwY*A{g`Q&_ZCoe^K95Uwj+w*93&uWrf?%94P85v&@Fc`nR5gPspAD;qALof2uv7O*BF?tr^BBbG_ub01<{CiVKZ+hT37@d@MV$0+ z5@AP=E)aI~NaetIicy!z#CY^p*|8U>x-hSm9rJohF&=hfM1JJJJ>t9@9Q0517CUwS zFUGsDGTudkHOjyj3q~w6iu{O0-8&05o3JBJ#h4S9h*;>YvZD`|3OnZDWr8u!E*A{D zhRA{Wa)q#?W>*TvIZlm7%@iZ|tBTKfeF3}LSbWd*8Zka~x>opsU6=e^FZ@uOWlC*g zJg^&*pBsfA>VH#V(GTqA%qc8picZ#7{@y`K@|v2U&+>{yr23OmjD?CD>?M7eC~L&irwadF|YdxM()ZdO!;KaZ!COL zmw|#Y2L}koHO7mgX0Vq8V?C%i*;f>!4~nslwi0ntPgOIR^2wZdS>%A-ATb_s?jjg* zZYLP88`O8}>Cq0t4)%(Olm6`_?C8;}!j2xP92ieA>M~f2M{kuKdzPvT^IF+4uZM{7 zuzOA9M-JNxJNl=3i=DbR0`Grr@VRaC>2C(G4J)kE&lB4*Uk(y+!d|cB_l98Dy;)&z z3C2GEwqVry9l3$(}spXeN7x^opOKra{y2#%MU26Sp(MA5Q=<HaNriSM7J`)@_JY0{SQwlO)s2BBt?B;84cE=J*<5hgL&gUq}EdkhR>-gY-+)%-86zxt7!$JHq!}4 zEv7Gwb>cY?=M2W;9yVh|cc!E}v(TkBvlLzAS%oe!%$EGlUeTQ+>CP#1spVWn7kO@> zOKs;Vy2$eiU1~jF(M4_)UA2eJZ!GR%3z!`DunxtR-o!m@A!A*`d|R;SG8f|>ws7(j z_pn8c#XW41q#O4z?&U#m7E5Esx#Ra0X>Cm{F6@{KO9;kzOIFxYg3;Th1*4bC2uAM) zeB~Ln{2T@AX4zub-c{`IIx${%W8PP;RlAw&eMRnSviE>%%PuB+50E>X?7bdlKYnyF z*?Wz=yvf~6re+;Y_I?xFawdD+)7Wo6JT`RKGUhU4tyyGj*C;Y$tzKkeU(Mv0TMv_C zZmXIcb6dsanA^%G$J|yjIp(&a$uYP8nH+Ol!DP=()qQUR_Q~%K72O?^?oL9NTJBtQk#`Zg)OOdRi#%B9QtRD{ zF7ob0SDn-MFc#10dz$Qf5aSFfw)7^R(}x-BYP!^LXwhX3$8-8#$&c@AecB6R8Dp-J~Jq03x3Jn0@GbXm7YCS9+0Yv+?A)A^+Px8A9i*DBQasN#otIP!gz|2`{^Q1^1H8+ zySi@&#5tQ_BZ}R$f-%p|7Bz#NBiN|oQ^m<^EXC-9Vw~${7I9KfRWq3K z*(h$jG3wGm zj7M*k-SARtRTt*9vSVH^EXKp`0+Am%%pvUPpXx1kvx^>4Pjz1f>Y^BNjuCOPw?19j zTbC8FaK7B&3omTT^AWLh6L#d+Loo7F<8ke<(l4!#ccI7+yI~?u*j^QOh!~F?E~<=o zrLbeX11jSkC>Y~iEXIRfA{gTxEqua`5e%P~3ZJmcD(vzKTi~0C`-jg3D{LXbSVvcg z@nBb0*j0j2mvzLtCC;mb9qgJ4yH>D=yc~91R~WBB{Wy7#Q1IRZB-JRUPj?2Dm zelG|4a+7_1D!(@-zc&@X@ihr+Qhj%pb*>oW-zT-XDJs=o99~6x1k%t81 zHQd9A^$7Vrl34tnPyS9$v{SWayhoEyACCF*m|*1hxM0-m3Bh>1`=nsZv!?`O4LmIv z<2{quCZT4}Ce|ykT(gY>Q*&ux;m;+X{xmgvUNG!lsIV6+?4`u|hB#|EWaehn?B(Rs z*%1F_$18#{UtSe#*JiV(*90T?*DLIe3VTy9_LsL3>sR9Nhn#!t``gLR^~vv@#MX9L zj>pCR!nx$#WalQ#iT4uo_ft9EPpl)r9y4=<&kvGaWFIEh#bG&=PuAZ@$<9p}?_espV-Hb`dVVT=()Yv0d1Nimu`u;L7-IQM_#_tfeI0V2M!YA8&)RtEdoYZr z7`cBg;v~N#Mef*rA?(QEOTijmMmxR|j5(pcUxH8dT^AVNiD=V2#Q5JZ|4u~iA9eBX zrDR5a$|p?u9ay^ODzZ;Pm@B66F_|v2NgdOZV5hwk7K-kfv?}Z&b zQaLc5V${XIQUI$a3TjqVA*S_?uTAm+>MSZ@S{DxKXQ{&NxpF|GK!Jh?Vp8X;i zcJGPaGGBfbcJ%Kz!8lK=@yJ~oBv>W$nz8Vwkn)dPg{jcM7o{sJ;4JGh@sq zd}8;9*sG}f?81)z`RB{?nxY3jXEI~Fe}o<5)$(f+xl`{s#CYtRlL$N3<($Hfb6Za_;d3#;=)>ZIQR^iHBbFrvqt?ofIj`bmo-HMO!j=|{9F`G`98}G)TUPXl zIF}VZnG?$iMy=I&^r*WS57x2bb9uqA>r`Q#1!E8CA{hH_SHaj@y9vge=w4X7-XYEv zjK$Bv{;#6DV$xkn=u+>Mi!SmiLYEj;O@4b+bXQBds|#Ifxkk}NUQ_5&+qH@=^4da| zTK6ov$i0fL`W);!#yZ*=AKl(2FK;sAtXpj9O&4R-rH`?$CR4-piZ0hb-GlD>$xkO^ z&OvvBq#Hj6+soEMx1ifM=^FYkzJA8yHFy7t?nX&>W1&l)n z?9pAs_kK8+C`M2GyA3%`dfP)V;@qJU=Z=CA=T3r=-&P_&&LukwJI*D$2*%t{Gcghbe`xzj}>s5{l|MiA5OMdykGRpzqgV1D(1l{ zrW-#$cWTK4zix@&t`@87B6lmXM_v4$mi(Sx{KhrP-@?>q`^i%=#y>;&v zF!kBuL7u1`{l4(uE$0TYsLwz%Hx?7`36dYU5UKX`qK>$Jv50$EN$;HCo0E zyQeb8?ya!7ziHj)JQcQ;h=n?xBXWS9TVdx3Mt)0HVp*!fUJ`MR@s#aY`hEZ1vYyO? z^GhtTpRfyz`5HkUBjSc#Sn+dF@^i7!WvyLObdfJDx^eC^=4HYU?DFCV`HJMn&sChS zu1svOG5_)7sx*erU*_=D#{9g-x$TaHz_GIT~Tr1om7<2N@#5N7%<=i(8Onr8aedn&^ z)6I4*ws#B0dG(&e)(_+Do_genn%$f1oH0M|6O1_TPppp{mLq@8G&0p&@_Qipba&+U zU}Ebsm_63U%YEn-?9^w0IZr;6#_JyJ-WS)1jlv&JcK$Rqdqgni(xZv_dvP3jyv_i7 z%*^+D4p=-lurEJmnUyjWA}{dB6+Ga@JQeO55W ze@?JQ$?18)SWoIRF8EZRw}JgD=02=Rj7LBHvoX28n0WBQBPGfnDWV-c&XxZa>0mmI>Csuy$jG3q&k@JVl#9cx_Gg?X*)nAfw2@vwVUFmjke zF#4x@`=6cY5#zlkd}8;yV9c{O1S6Jt1jDELj0J4Lirqqjk%O{>DLdxUo5CmTEy0*C zYCM=4kGd=&e6o&~5hziL1N%Jr z`9k=iHa&$O*q6!ASHchV|62HgeUtorEBvs>DnGF7=R4tt-YGw@?B{#ohx6SUg~fis z+%ER(2jPb~@uTnq`ziVHy60t^ zKC#RLMboeu43ZRqXZ>jCnm+FmhKu zVag|Seh=Z3IQJHeIk=}_T<=UNY6hE1FxG>LlYK=o`k)x=Xrzdfda9biluzcw)FKD$ z_7UR|=OKa-=YE25-KqX=OOFl^cCcwgob>M?VMmXq6?XJU<-mA~QJ2HSc=T4;v1h5e zFt3#z^ZF<;9(L1-{K#QnVMqT|Z?RK_kQb}Nz3=OVE0wn*<)MU`M%2DI=#q& zTF+2eyf%W(Xe|CcYNm>xnJa!~Nq%M(y7X|iqKiCx(T(S4;+sSGfz4U`AkQWIux92K zjG8a>eapJK1Y5YmW*^tO-8_}?<}EDtgYo7w7O!g>6+iP=SceK*Krm{!U|}&f*h0qQ zTv=H7A&*5WY*8@=c8eK{`7fUQ_@2bRv{m{Wiuam+FOlrrKVxj2>^whmyi%<9hH*cZ zOg{ZN*85V4#lQdM`#G%dhULhA3-PCi6t;Bo*}FXN59;q#*4VKmIv( z4&q`j?h3HS%p4KZ>cywpH`Xu~_l-5v_|eZ=6+deyKRtym^R-veMP8@ql2;saUEznk zdKW*)>j^*9S^fP6JH?oLeS}Zu-k*PYq!W7O5=^~+Of2gQpTwfx=VtEBF7`is*2bGh z*fE}Bga> zYY`PEuPGFx4~lVa^Y;#O4%AcC45oZGN}?N!9I#XGbr7e&mzHB8PJiz#Gp>cydqDKa z-%reTuuVjq^v~a0%y#r>fUu)SDhI|>jJo_S#-q2&j(uO%g?X*)nAiUPV$Kh{O$8%| zUqqbrPxThNZG|26RPQTM7sZHkGZ80y>s00cYw;{OVEc$zb}LHr{owy-c|IZ*_4=Fq zjuk%1FOTQFqko%=9Jr3(qQbT;EMAXbx0SIaOFsUvb;Zv%6*jQK1_?%;wk<5i2HVb9 zyjI#?_#uxSDr`qF26j6ci}~*?{A^w7vx{KN=UppoaE0wA7_s^5-#ot^VTon`WaqKt^MJ(SYs#9>1CyQmnY3g!F zV$tr<#N3Yj4ik)64i}8^j;OFB1*2w16&B|*=ZRy+vkBjAK3eGhw~rTH;yF&}!cHi< zyiPs7=*DvneoqvPxqecGom^PV3%gT{#Wi|r#m{LKc6wn`cmwT-J)ZpMrl!8lPmllY za=#!?W^(@^w=dY@}7*84o6%X$5LV;&Fk1wxnkxvu5xy2!mt4Pu`78I$~8Sp3F0M?KZ+FY2Wj z<6k6va_w}C_&-r#>iTbxH_(m_KlpEbU8Sy5SvPG(Eab;{s!kWDI$e_d46|tb$B#=B z^XHhemsQy16?TPS%p<=(%e7#3dyIqV2yQ9MH z6pWhrbz5t3-c@0DSJ*v?`T34>-o1sj^Th1%_0F32zhC(4DqdT;9RL5UZA^~qa7&YY zZL_}bOXC{)kKLbG{CmXF<|o%+pc~uqK(h1asNI8!#cMfrKXR@E9!hrZlkpxF47*1X zi#e#jn=m&XO?K<}faOqj#PV3Oa}#2DTrl1%ctSAFBTov(x_C-3=FQWBk>4|d@xIKn zg7MzVbAoZ+d|ohoz91N}yjWo`RoKfF_KIN4LBFQUbC+0Nt+3Yw!{_T2_C{f9|9{h1 z-2dMaeu(Yu3VWx*-W80v{dz0cf_e8|VKHvl`^Nlw3i*SIpARd3K1zN*7P`!@Pl_(` zr$yH@;2@sQgdf=F#Siiq$&asP|MBC?#QZtd-d7d&b%lLXVGoLP8RI=vVGj$&y!lp) z2m7wVz88#nD-Vo;x74~C={UjKD_*pPw`CVKaz{XbC4C4PfActSXc(7k9>^H%P zWpdFY*c26ZSc=o+?-q;qnO^75I=Cg*=VMdc8#pXSem|ry zz#jA8-yPAjami=DV5h!6O-_F#J2#{DeKK}PXmDiOWlFvTD?w)jxbi{dRfn?{5*S`xU*4<$_)E?d_Y@z?t+j(?rQ3PRF zWg{;54fn86l%R4if{0=O1ruUMObCbyA|}L~Fo6jaB^W?4XABrn`3daVxbaW8a%yfp zPM!C>r8`4-uI`$dZ_a$uFAg*BCmKe4o+NB6<;92Uo7G^k<|x`Pjx$G_ z8*}S^V~Bq4H-8x46d6P4gPq)9rx=F59Su9(Z)DKV{bmI*Kh@}iO>3~}hM_O^V$B)I zoY7#XHrQ!~q3`qto0+Wc%k&%VSlu(nj2!Hq)nI2d*lfelAA5H7Qyho!UNQEK4*kde zZp49m-xe}?|5cUA8lNM!#PuWBkY@_3=csdq)%`}Ox@LcA(S!aKk zmvzq1&dxgL>GQMB*v`p1&t;z1*N5|~QH<(7u?ru~Zm>Z3_K11+2F$&e1IBwY&=)e< zC(ai;E4!e<7Btwg=6MwK9oJwJ3}bC9H2Po{HrPdm;g?vOb#2tKFKV#G4R*0%b3<@f zE)f>7L_A$8tlp1x_c^(jd71PJJ+R9SBfgg;t9usLVM}{9e2096Y0G?GnQhT|m1zsR zCfl;6uFkeDF0VC=dHcsaOA2$(*uva1uZ-h5qYrj{gDo=*eJ7ZBhPh|9VEPPI%;`w{ z>&gwHZ$)=3{-P`w7RR`+eWPK-!cB%T_L~i3{%$dhHF>LH+=E_W*rqVHE4LX&j#z0J zGFKUfEvpT~*IUeYyfF9uFO2VFYwS)QAHp%Wi@xwPzFs3Ny0GrH8a5C}r-O{hhTyGe$bWgI0(IE?VZs^&Ns~^ldKD{^V$ht4< z?5z)Eom_o?*2&cyyLz0*)BmBY!zUdLgN9rlF+dz}->=32dN%cXhJN*$#`$n0)O6)x z=@-ZF#UqB{hew5tMaL|vuU^wAtJgHjLLB?~W1_FJ#|^{ICz83`^Hk4<-1DSqi!Ywe zw%GHGX$#w&ZJFa|vu&L>?|44-|Ltj^m=bneRd-ZVTPeBiG*I_wEiV{LzKbPyX~7>0kpGz{CmGK}$l zZ5ZqD8^c)J-x|i+{?0ID9xx1BzBi10^#@@!XS?~tFF#7Z@HPJYNm$LxDzoNgWi>A= z3w^AQgQ71OcK$4^t`EH*RM&^HD)SeSS=q0`>iTdpnWNvNUzK^tFkfMI|c+-8yxclR_&hzoH~&P5W4<1TSG;QBrH z+~-rZ_wAQjbNK7tb=HD@*829dzx7u2R99F1_J#mC?YXbs_{vwj;uZ6Uy$?I=q5t~T z8yCv*%4z4!o_YBqSp3@hoSZYP!PkZJ_1xm_Gp)Gq ze_rtgvz}4hlf0z(0$I;2?s?5F?x*lv`tz=T)}ooz$1c|H>(7vjhc)=R_~7fDoBHhX zGFxI~Hnq%#FU!0iw^V!Fyq?=KnN2U>XX$p|{^!>}W7(1KQ)kO(R)6N@UgZ0%FfyyJ z^NJ(0`aM?~nbqI#$|JM-yIf^tR-gB(BeVLr)kbFZ_q6)Rtp1+X7@5`g<(eb2`W{+q zWHxKy!!K)(%!%#4BF0 zXnCC1Y%85_tbP`@Z?kzl zf4_tJ{CR!;d@S_v`uul{*XO^(;Eof|zkarMYO{Gg55Lp;Ja~N`d@S_v`aE`t*XObG z;EwZt9>d@J#H?OZc59ED*K_n9sn3zu=g7xG4Rk^^FY1OC@4fo7^7^y#vCzY3bV=_Q>UHISc>US;AKZ81 zXP-9k;g^F(X5(e%z3FGD&$WM!56`6D?tNQ-23~&#J{EfTd0o=;{GmjjZjSqbI#rxdP`_-Yo;+1L0$zi@m_xZDcm@m-izF-y* zbMN#1Ugyq-5WYyGyVtkB#^zq%{u=Ye8>zwLmq3}?n2D>Z*)H(3yArFjqa|!!|U&GesCP( z{;qKl?ivZrV6TBG~vSwPIsXmmd_3yAqyjnu%WrXl8MH@csb1;qT^M)zr1 zK+Ml;bU!}}i1`JL?$fh?m|xiFeo+<>^NSnZXJi2}zogN9W)=|hOB>xU%K~D4d87NR zEFk8y8{MzS0%Cq;qx)4^K+LahbiXDGi21dR?$>1jF~7di{e~e^_w*0*zns%_upj!G5>v|`yaA^nE$cS{ZCmy%>Ufz z{+BEu=6`K;|63Li^S?K`|04^CYwG34Ag=BIY4D@NvyJZ0WdSjNzQK?3 zC1>SV4tnCcLS47(Y<%FuFWU@i-j6%=MP{z`^?Obo`0&f=gPQoBu8TePfcm(8{d4P? zzLcL~eBi?`I}U2%drm8}vqomrwOoeJYuTAO5Yr@Au53jGuh4K2D44>zza|Ta`OS^)ugwBt z{<=o@*JlASe?z1D8?%6zzp2sv%~?Rq-_q#*)+`|AZ)G03Am;CBbboIa5cBsny1zdQi1`N^-9MNG#QZ~z?jOzqV*Zgv z_m5@)G5=Vj`*;=*^IIC-Kb{4|{1c7tpUeVceru!qr?P;Uf4b5AGg&~)Kila3xhx>& zpKo;kLKYD7FE+Y=DGP}CmmA%0%K~Ell}7ikW&tt(TBG~dvw)a?qtX4FSwPIc)#(21 zEFk9JX>|W?77+9AHM)O43yAp-8r^@G1;qSEjqa`mTa(1+0OG!rpM~M)p}!uJpM{tH zXJLw8b~2}7eL|@FB6mK3@OnnzuW6(P^Lj=wuV(~a&j`Gp5qLc#@Onnz^^Cyl8G+X` z0{-gP*tVi{^D-__eEa!MPBzsUiU>__eEa!MPBzsUiU>_ z_eEa!MPBzsep@3o`1$L;XkPb4{bzkIlU*vUP zO!31ea~jqs zhk8cf^^Cyl8G+X`0luO9GXk$?1YXYwyq*zwy)N*2 zM&R{~!0Q=-*E0gIX9WJzwTTh9pQ^^Cy3(?|{OU(X2U^^Cyl8G+X`0Jo0%E>cqx<4nK+KnDbYC(Hi1|{D?n`F@F<++9ec3D^ z=F2s@FP{a(e1%5$6|;btuhi(iauyKtRT|w_%>rV+TBG~wSwPI!XmnpQ3yAqzjqYn_ z0Wn{v(S2Qumeca$^M0HAuh(!FzJ8b-zf`-`Ob~*yJP_| z-?hrV6Sfl&lSwPH>Xmmd^3yAqqjqXQh0Wm+O(f!ygAm+z4 zx*wkf#QcOt_Y<>#n4i?>?ly=%rP2M=EFk8mHM*ak1;qS}M)xzbfS8}vNDX{y8e)ES zqx(5oK+Ml=bf1<5#QeNQ_w%!Wm|xK7K0OPF`Gt+{7i9r4zqrwTMivnBOB&s0W&tt3 zw9);tEFk8WH@eTt0%AV9(fx`nAm&##x?hzA#Qf?;_iM6%m|xrIeq9z2^XnVkZ^!~- zeq*EiO<6$9Zx3TZycapd`}B?m_g=oU(fzJ0Am%@AbpJ^f5c9hm-S5c)V*b-c_n&0} zG5>j^`!BM9nE$fT{a0B)%zxeJes2~K^ZOdz@6Q5a{y?MqgIPe#A8K@eI17mRBaQBl zW&tsOtkM1PEFk7jG`c^T1;qTRM)#+)fS5niNKO9pV|fk!)9`v69iDA;e=ZA%`ST5K z@9@{uuN?FX|2@~?HELqkuLxi3O2m9xq53i9E&%<@J=d!_KV%GNyH@jZepN;kXU=}H_*Z9DPUz`KX79HyU z`_}-}ep7!oUVk<|7J4{(TyO1DJ+8U-rJgIZsJ|cU`dYZg*XKR-!Hw-&J@41TwYxsf zzSm#J;cGwfIQQ)Po@;u2T)*9|!)yD*p2zsWhhKaa{XQ2N>fo!-o4g?`I#>Fi!`sclMWHz-wDE+pPS^b^A{Y7T)7@5^;#5+f3{=9xoc-P2`_MhY3 zBeVKF-!n4nk8$&RM`q&#AAY&;MP@!L-b;%Q`i1YExw);rcX)m8@UhUt>wD*Y@%kJ; zP~5-s`QhR#f9u|zKhk_^I`i*-sCnBve)Pe3eebUp(qJ+1J?pjZ z%k!A^_v39NGuj{fl{TB#p5pUf-ygibKloVa;q`s`wRn9$e06Zgc|ZTJkId?OVo}=y)decfjlKfRBZq@hs~5>4)+DGrP|ZLj9e5-;MdRUpeR(-uFi%v+;Hx>*3~~jm-QGdi;Hb`unMm z`^7wFV}sjNvlvCzZoYj9tD zxF-znaeu0Q?(U7(-_frJcbj;g^>uk*WLAGi501=c41Du_Z|NIkcmIf<1T2A z_4h~W>si0&@b8nn@Nw;7gTHU$ar1glyOz|~lh@aikA)sy-_!n{OMSf_AKY=`d5jNy z_+_}BnAO*P`1ffhW_=&_$fpJmpP1FxdHDBlCT8^-@XW~L`fKdw-;B(r#{2X6?Z}Mw z>&oxiY+kRWYe#)8d3`PUSm@#PwfsZ8zLviq+;QS(8z1=ai))MD>2M!>;r}-IuT17^ zxEG(l(S3m|Am$4;y8HS1&uDHxfpDMA{`usKHd2GR*P*`_=Fa5(T+Wwhqz3o*I`!Ae ze5poiFkdrV+TBG~w zSwPI!XmnpQ3yAqzjqYn_0WlwIqz1lD8e+a)qxrV+U8DQV0WsgD(cKxMpCin7YxMcUIgRe;W&tsu*64m-77+9E8{IF+0%GnA2;mntx?hw9#Qfq$_ZeA0 z%r9wlpP2>3{L)7E%d&u&U*70GD+`GE>_+!1vVfRh+30>%77+8R8{MzT0%Cq`qx*GP zK+LaibiW}Bi203;?l)xtF~2>G1^H*-kbi_b8vO6!&PMmUvVfTXxY7M5SwPJ1ZgjsV z3yAqo8{L1F1;qU4jqbn50%HElM)zN30Wtq|qx-#CK+NxJbiY3fi1`DJ?hj@GF@LDh z{oyPi=8rVGKbi%^{IN#&$FqQ#Khfy^WEK$fryAX#&H`foOd~b;Z$$hi4f#v>ZKM0| zvVfTXzR~>;SwPJH*y#SJEFk88Zgh7WME|wX{cl-7%>4`y?mY_eSnpwox%W7P|GU9G zm;W@l=lN`-`*T@9%%5*?KWqK#rvABgt3RjCW(|D!#h+Pct|2}C{5qT72kGJSm^io3 zC^OHiKF)Qk=YJk7CLTAl57Ph3$jtSs=YJh6CLY(X<8J=nMrPv!AAb4T$gKYC|NF?S zK9~PvWLAHM|8rz^{=kP{{;!eQv@-jzFEabT=P?^=kNt1+nAsZ#d8MxtuS@;&Uil)k z`9@~rgKswc^1VULxoP;?z3{KQmdxYk%N)YzPeb?uX$W614dDx=A$;L9gfEhY@I})Q zzE~Q<7f(a@5@`ruQV(qCOJxl)Upfup%cLQE*))VNmxl1=(-6Kw8p2mhL-Yi z2;V3T;TxwRe3LYUZ<>bi&C(FQc^bmENJIFRX$apc4dGj-A$*%Ogm0UM@a@tNzI__P zcSu9{j%f(rDGlK}ry+coG=%S(hVb3e5WagF!uLo+_?~G9-zyE_Q_>K=cN)U?NkjO) zX$aph4dMHzA^dE}@zbOsjx2MNQvP5&Y$>Q-l(vZdCccvlyt~7-II1S-HNkjPE zX$Ze34dFjcL-^0q5dQNtg#RK9;lE5n_^;9s{_8Y^-B@N+!O+)zK(h&alG=%>n4Ot}q&ospRUug*ccN(%# z{6A^Pg7Ih55dK^m!kz%*d=?6O5M~nN8~tN-sAuo0{ADW0xP9onIce z!pLlPX0zIS#gUnt&1I#L+4$hr;g{j(eBS??%vc^bU+xCs^QR$vfi#3Kn1=9$(h$CI z8p0PzL-?X;2wyA>;ftpse2FxKFPVn$rP2_-bQ;2!NkjOuX$W5~4dKhDA$)~2gs+%} z@RiaKzH%DES4l(os%Z#cEe+wTry+ceG=#62hVZr05WaRA!pC$FzD^p#*GohA`e_K? zAPwOgrXhTzG=y)QhVV_&5WZ;|!Z%Ap_~vN{-y#j+Tc#m=t2Bgfordsj(h$CF8p5|r zL-_V-2;U(M;X9@we5W*o@0^D4UD6P~YZ}6LOGEhXX$apV4dHvHA$+eigilFB_}*y< z-zN>>`=%j$zchsJpN8-Q(hz=N8p01sL-@gI2tOnZ;fJOn{IE2HAD)KrBhnClWE#Sc zN<;Y3X$U_i4dKV8A^f;Bgdd-V@DtJyeqtKJPfA1h$!Q2bB@N-HrXl>aG=!g?hVV1e z5PoJF!p}-W_|!CnpPh#AbJ7rgZW_X;r6K&hG=!g@hVTp05I#K(;TNVM{Gv33Uz~>U z8EFW=Bn{y+(-3}X8p1D2L-^%s2%nXP@Y!hyzakCcSEeESsx*XOordsh(hz=a8p5wj zL-_S+2)`i>;Wwrs{H8R7->ylrc5}DMTJby5kTv6XrXl>UG=%>+4dFjYL-^fk2)`!{ z;Xh47_|MW1{_`}1{~`_Hzf42;uhJ0y>okPln}+cF(hz=s8p0n)L->Pf2!ALI;SZ-F z{E;+-KbnT{$I=k~cpAc=NJIFOX$XHR4dG9xA^e#%WR3W5(va2SzfD8<@6r(d`!t0A zAr0YwOhfpe(h&aVG=%>p4dH)HL-^m)5dQZxg#RNASv~&GG{pR0X$b#!8nSBqKWWG+ z@n_Qz{#+WupHD+pj{9qYWmGb!VeMy{+vxtPEFk8uZghW577+8B8{J=<1;qSyjqb0{ z0%HD#M)x;n0Wp75qx+k)fSA9f(fzGiK+NCP=@Zx#^q_cgk|KMRQY2O8Z!m<7cALyhhq&H`fokw*8AW&tt( zSfl%R77+7W8r?sh1;qRljqab!0%Cq^qx+|_fS7-}(fu=7K+Hee=>EAZAm*QMbpJvY z5c4lKx_>DPi20Wr-EYeRV*Zsz_pfFFG5=bl``5F8n17?u{hL`p%)iy>{_QLv=HF>_ z|85o#^Y1mfe?JR|`41Z1f0zZt{6~%Mt_53@#ODCwzLTGY;pd^h9+RJim;PsAieGjz zr(yl~Q1?add;sC~jKE*hNDb!oj9^~R2)v#Vcs(QVH#Aa%?^n+V=Jkxg-`q$I?qAOc z=Jkxg>luO9YXPs<0$%q;UiU>__eEa!MPBzsUiU>__eEa!MPBzsUiU>__eEa!MPBzs zUiU>l-bf8zzq&7)*J}Z<`y#LVBCq=*ulpje`y#LVBL8e7HF!RCUo@}#BL6}oHMoD> z7tQOw$m_nyZ)>CmkFWcpdEFOz-4}V?7kS+mdEFOz-4}V?7kS+mdEFOz-4}V?7kS+m zdEFOz-50riact0&*MgV(wV=Ki#`B(lc%MvO7bgDujb4wJ!*yXweomRwu>M!5X9Ql) z2)v#Vcs(QVdPd;&jKJ#|f!8wvuV(~a&j`Gp5%`-MsloHFX9V+lM&NI2qz3n|*9G%> zM&R{~!0Q=-*E0gIX9Ql)2)v#Vcs(QVdPd;&jKJ#|f!8wvuV(~a&j`Gp5%_o`HF!Ph zwZOcd5qLc#@Onnz^^Cyl8G+X`0?Te|H}5+A8)Ptb-}es+^65qY`#WnaNqgU5c35Z-51OP zV!lwL`@&g3%ok~NU(}*y-8>QAZ*%{}`Y>(I`)%fnH&TP|zeE~hzT{vv{NnMv-)6p4 z!#()Yjqb~20Wn{;(S5lrAm+lqx(i# zK+HF8bl)Tki20_C?we%+G2gt=eTyt0=36$pZ9 z$O2-%W25^{SwPHpZgk%z3yAryjqbZ;0Wsgb(S46BAm)2Ex;sDgGlcn+Mrv^Xz0(l$ zeHz{O%>rV+U!(i}SwPGWXmmd?3yAqajqV3$0Wm+M(f!aYAm)cPx*whe#Qcax_an1_ zm><>XesmTP^J5y_kIe#Neq5vb@mWC3PiS;MF$;+KNsaC&X8|!krP2M=EFk8mHM*ak z1;qS}M)xzbfS8}vNDX{y8e)ESqx(5oK+Ml=bf1<5#QeNQ_w%!Wm|xK7K0OPF`Gt+{ z7i9r4zqrwTMivnBOB&s0W&tt3w9);tEFk8WH@eTt0%AV9(fx`nAm&##x?hzA#Qf?; z_iM6%m|xrIeq9z2^XnVkZ^!~-eq*EiO<6$9Zy&6NU)Fi$3;(&Ff3|tQ&p+aKG(6(( z;m$_)yRv|o|G3fpCs{zu?{0LzCku%APaEBTmIcK8=Z)^a$O2;i%SQKKWdSk&b));e zSwPJ1YjnRq3yApxjqVR-0Wp85(f#2pAm)!Wx<8r)#Qd>F_s6qE4X zAm)As2=^X^c&zs@#N2xx!vEdip38q4-1B_4(fzqBAm-0ExS!QuqvvBZh`qaK{lSg% zo^3EPo8BLk-f(1gNttalGMidv8;{J+EVE5UX7zh+Ix?G99=F-Z%(by+^O0G7+!iCV z8T~=&Ek|ZE%WSKWS^b_{kIc?5kK1NsR-fm#^O)I#WV?CHE*_3vvHd(|ZS;yAUSzgo zo6XnmJk}qS-YK(ni|^d-yG*-nP8zasW^>aJ{;D*Dzd8-!uSrAr&1ndKZ5qN~mxl1y zry=|eX$XH~8p7X{hVVD1A^a_A2!CrD!rzvLaG#|qDw)%;Hk+G<@R#~)@%H?j5I^rb z(h&a6G=#q^4dL%jL->2r5dPjYgugEh;qOmF_y^Ju{=qbae<%&%A5KH~N74}f(KLjA zEDhn~X$Ze14dEY8L-;4s5dO(Dgx{Km@K2>7{L^U&|4bUfKbwZ|&!r*!^JxhGLK?!q zn1=8#r6K&wX$Ze94dGu&L-<$I5dO6^gnvB^;onF@_&3uK{;f2Ge>)A~-$_IGcheC5 zy)=Y>KMmnONJID!(-8inG=#erY^@TX1Blt=XJPWQ@N)YsOc}iUU$`&&cY5dMb_lQg zBClrz?tB2@H>V-Io)LIGBk+1g;Ps5a>luO9GXk$?1YXYwyq*zwJtOdXE#UQ9!0W!q z>%PeAzR2so$m_ny>%PeAzR2so$m_ny>%PeAzR2so$m_ny>%Pdx^Lv8uThb6-_eEa! zMPBzserp~F;dNi+bzkIlU*vUPp;I2@VYPZ zx-as&FY>xC^13hbx-as&FY>xC^13hbx-as&FY>xC^13hbx-WA3;+RV2G_1|$rXhUt zS}<`&_IkY3t_AhIFrMEz#Ls$58p0>93lryMug6RAx-ez%+kWA7!GGsvZf=M0dR^f4 zjKJ#|f#00RL3lkQ@Onnz^^Cyl8G+X`0luO9GXk&I z1zyhxyq*zwJtOdXM&R{~!0Q=-*E0gIX9Ql)2)v#Vcs(QVdPd;m`8`4SEolg^*8*P8 z2)v#V_^o*ygx50yuV(~a&j`Gp5qLc#@Onnz^^Cyl8G+X`0luO9GXk$?1YXYwyq*zwJtOdXM&R{~!0Q=-*E0gIX9R9v*h?jI8rEiW(-1y+ zEtohXcs*Wf*MeK}dxrS=CO`WV=K!zAOYyUxA8U~D*2>T7y5Q#(+V5vJUmU{cPeb?u zX$W614dDx=A$;L9gfEhY@I})QzE~Q<7f(a@5@`ruQV(qCOJxl)Upfup%cLQE*))VN zmxl1=(-6Kw8p2mhL-v+gRi^f{_AB0@%Z)A5WYbg!Z%Dq_(o|6-#887o1`Ip(=>!{mWJ@n(-6Kz z8p5|sL-~gzuJy z@ZHl8zDF9u_e?{${lA|9_>?$=@12J5ebNxVZyLh)OGEhnX$U_c4dDl-A^e~;gdd!S z@I%rNerOuP4@*P%;b{mzA`Rh3rXl>OG=v|WhVWz35PobL!jDTs`0;57KOqg_C#E6% zq%?$|oQCjI(hz=X8p2OYL-^@w2tOkY;b*2H{H!#DPfbJk*=YzrCk^4}rXhS<8p6*@ zL-_e=2)`f=;nUL)eqkEIFG@rB#c2qik%sU~(hxo~4dIujA^fs5gkPS9@L6dHpPh#A zE7A~tWg5b-N<;Y7X$Ze24dK_OA^f^DgkPVA@Eg()eq$QKZ%RY>?ddU+O#ZiCU--{? z`TIfl=XHVb2htG!U>d?7N<;X=X$XHL4dIWbA^fp4gg>5!@F&s`{$v`$pGrgc(`g8Q zCJmYVZ@s?opYvKUug|l2T_F6qG=x8&hAa@DJUe?0hS!p1(-7B|<X=6NJIFF zX$W5_4dE-NA$*lIgs+;0@YT`~zIqzM*GNP7nrR4MD-Gdmry+bS4cR@uP8wpq?r=Hy zx@&yBtRUv=ry+cUG=y)MhVYHj5WaC5!Z%4n_@-$H-z*K`o2Masi!_99nTGJK(h$CN z8p5|pL-@972;VLZ;oGMne1|lI@0f=0ozf7#a~i^TNkjOqX$apf4dJ_|A$*TCgzuS# z@V(LyJ|zv|d#53MpEQK;n}+cH(h$CX8p01qL->Jd2tOzd;RmN7{E#$+ADV{n!_p9b zcpAcwNJIFMX$U_m4dF+pA^ey$gddxR@Z-`Deta6jPe?=fiD?KwDGlK#ry=~5G=!g; zhVawU5Po_Z!p}%U_?c-4KPwI4Q_~QBb{fLZNkjO#X$YT|hVb*!5Pp6d!Y@ce`1CY{ zUzmpQi_#E&aT>yBq#^v0G-S8<%rwOO(lmr$mWJ@l(-1x@4dJuX5Pn4(!mmt2_*H2L zzd8-!*Q6o*+BAe;mxl1`(-3|`8p3Z(L-#4dH)EL-?Q55dN1mg#R@S;eSg*_}|kI{*N@|f7Jg^!qHpO=R4 z^V1N1K^nrRry=~pG=yK2hVYBi5I!Rf;g_T#eCBXD_&R^{e%s{#=h);~`rf=VApSn> zeQ5~4KMmmzq#^vlG=x8thVX~e5dKIS!XHgT_+x1Ze>@H0PoyFI$uxvNm4@)A(-8hl z8Z!C!q5r@6`_Oq=2kF;Mvd7@YiGSa?Uz5!C9GUg&quE{~v)Q?AmO(OQWY({{W_yp! zTvNvkl6^*IZX5F@*|*L5nRrfagZTetn45<1SEV8R)oBQSO&Y>)PDA)>(-8i;G=#rC z4dHJ{L--rh5dNk#gugir;crPp_*>Hu{2s5dOY2gug!x;U7pt_y^Mv{-HF4e>e@{A4x;_N7E4g zu{4B_ry=~7G=zUV4dI_iL-;4t5PoYK!atRU@K2{9{4;3?|7;qluO9GXk$?1YXYwyq*zwJtOdXM&R{~!0Q=-*E0gI*8*Oz1-$Ny zyzYy__eEa!MPBzsUiU>__eEa!MPBzsUiU?QTb?I`yAJeg0k8Wa zulpje`y#LVBCq=*ulpje`y#LVBCq=*ulpje`y#LVBCq=*ulpjmFOI2XPQ%)4ZW_WT zuLTolWUt3d?OIUZ3*-5nL;S3_q#=Crx-fBG_IkV&uM1NKzwH-Z7sm3l&du!*Uat$h zo)LIGBk-H^I0&z21YXYwyq*zwJtOdXM&R{~!0Q=-*E0gIX9Ql)2)v#Vcs(QVdPd;& zy1?rhf!8wvuV(~a&j`Gp5qLc#@Onnz^^Cyl8G+X`03qngYbGr;Ps5a>luO9GXk$?1YXYwyq*zwJtOdXM&R{~!0Q=-*E0gIX9Qla z1-zaScs(QVdPd;&jKJ#|f!8wvuV(~a&j`Gp5qLc#@Onnz^^Cyn3wx<#PQ%)4ZW_WT zuLTol1h2d>?N<;X;X$U_g z4dI8TA^fm3gdd)U@FUU?eqeG=v|YhVT>85Po7B!cR&= z_{nJqKP3&}r=}tNv^0dDo`&!<(hz=T8p6*?L-^D*grA*;@N?1-er_7Vr==nMyflQL zpN8-Y(hxp94dEB2A^f5=gkPM7@EK_cza$OeGt&@$X&S;WOGEhOX$YT{hVa>G2)`l? z;a8?1{HipBU!8{VYtj&YZ5qO_OGEheX$Zd|4dFMYA^fH^gx{`7vUYQ~$y)I{(hz@t z(?Y!hfBH@O#q` zeqS2G?@vSc18E3Po*LJ=`@5t zlZLDj|4ka=e?$Fk8p3~)Xm!vBU0*eBI*fgILYpS4K4PyMspxBgkY{#ksiJx=lSiXRZawD^JXetmKC zY5C^$x$K|0*JNHl_xhRD=fdlA;bR`&!t2j{aJ>Fm2gU2_Kh_?n4jGxvEVDyLW>d@T zu#wrtWp?<;tiCQsjLhoyJaS~VYx$l>jm+xfjvkp^ULJSM$gDo@*pb=n^0?zhX7zg> zKQfzE9(TgXY-*XEIFDJMpEQqIe}A1kk6Hh=r@YAQ)HeHiyE7f!ls+$u`O*+Re;UFU zNJIF7X$be;?eAZ{a2&#Y&iZGK`@Hqf9QPjY?{mI*9Kx4KL%7dC{~Yk8;t;-c8p4-J zL-?|32wyG@;mfBXe1$ZGub77LmC_KtavH){NkjOmX$W5}4dJV&A$*NAgs+)~@U_wq zzIGbI$I=l0e_2f5&HHWV$1lwU z;{KPVA^h?*gwIMt`0O-c@Eg+* zep4F4Z%>brWLk5#$+_`6(vWlFccvlyt~7-II1S-HNkjPEX$Ze34dFjcL-^0q5dQNt zg#RK9;lE5n_^;9s{_8Y^-B@N+!O+)zK(h&alG=%>n4VfDMXBuMuuQY`JI}JHA{+~4DjQF!@2!AdO;m@Zb zr^oAGL$M2dNiJx!g@!u#I@Z<~WwvO!{lfTSBfjK_kBxZ$jJkRHi~Ly^|Nr!7eXTto zfArvwPy6>1=SxHQ{AmbZAPwOQrXk$x)n7Bda2&!HNkh2TvA>qw>)2mQzIYtMmqS+jHBMsqerXhT-G=#66hVZd8Wb)s)f8k$0czw3b>jL51r6GL#G=%SvhVUKJ5WZ6y z!go$X_%3M(-!%>4yQLv~_cVm>k%sU+(-6K_8p5ZfA$;#NgzuAv@O{$|zF!)`_fJFk z0ci+7Fb&}cr6K&_G=v|LhVVnv5Pn!1!VgbF_z`IcKQayBN2MYB=rn{MlZNnP(-3}K z8p4lHL-+}42tP3m;U}db{Nyx*pOS{~Q_~QBS{lMnPeb?_X$U_v4dG{{A$)2Y!p}}a z_&I3^KQ|5G)6$R`@$=FU^YhaXenA?-r>7zO!Zd_ml!ow&(-1x*4dIuhA$(>U!Y@rj z_+@DbzdQ}$v(gYgI}PDiq#^vuG=yK3hVZM?5PnS>!mmw3_;qOrzdjA&H>4r_#x#WA zl!oxz(_tFcS3k&A;c{Z;Lgg=*t@aNNz1>*jC;cF_PgPYR+ zyPxx=A$2eC@%6HTn6ICP@D0)szF``|H%deJ#%T!OBn{!4rXhT@ zG=y)ShVU)Y5WZy^!naC8_||C%-zE*=+omCWyEKGvpN8-q(h$C58p3xJd z2tOzd;RmN7{E#$+ADV{n!_p9bcpAcwNJIFMX$U_m4dF+pA^ey$gddxR@Z-`Deta6j zPe?=fiD?KwDGlK#ry=~5G=!g;hVawU5Po_Z!p}%U_?c-4KPwI4Q_~QBb{fLZNkjO# zX$YT|hVb*!5Pp6d!Y@ce`1CY{UzmpQi_#E&aT>yBq#^v0G=$GgL-?g>2)`^1;g_c& zd{!F5XQv_jiZp~@nTGJI(hz=i8p5whL-@662)`~3;n$}j{Dw4y-ToJz`4VfLkGY#Q)r6K&sX$b#G8p7{RL-;*u2>)ps!he>A@Smq4{1<5m|79A& zf0c&tU#B7b-ZX^Ymxl2B(-8hZ8p0n;L-<2!2!A*Y;g6&t{LwUoKbD5@$I}r0L>j`M zOhfonX$XHh4dKtEA+zGYNkcA+|27Tbze_{-@6!( z4dH)FL-^m*5dM!e8QUpfup%cLQE*))VNmxl1=(-6Kw8p2mhL-Rgm00C@Ga92zEv8+w@yR&HfadoHVxt1r6GL#G=%SvhVUKJ z5WZ6y!go$X_%3M(-!%>4yQLv~_cVm>k%sU+(-6K_8p5ZfA$;#NgzuAv@O{$|zF!)` z_fJFk0ci+7Fb&}cr6K&_G=v|LhVVnv5Pn!1!VgbF_z`IcKQayBN2MYB=rn{MlZNnP z(-3}K8p4lHL-+}42tP3m;U}db{Nyx*pOS{~Q_~QBS{lMnPeb?_X$U_v4dG{{A$)2Y z!p}}a_&I3^KQ|5G)6$T4#Lr7Z-X1?c4RQYq(hxp94dEB2A^f5=gkPM7@EK_cza$NL zTYP33^49pJX^8t@mWJ@l(-1x@4S7p^b{g{L;xm_f^{(%YUy<7(9)D#T!mmn0_|<6$ zza|ag*QO!-x-^7epN8-o(hz=Q8p3Z%L-_6KF_OHox!dFo@jKFx*T?TnL-<{32>)>! z!he#6@VnCxeoq?0f0~BypQR!E=V=K4MH<3?nTGIRr6K&+X$ZeJ4dM5tA^iR{gg=mm z@CVZn{!kjiA5KH~BWVbKG!5a8r6K(BG=x8qhVUoT5dKse!k*BvjLvD`$ zHVxswOGEhY(-8iLG=%>#4dH)EL-?Q55dN1mg#R@S;eSg*_}|kI{*N@|wef$ZA?E)| zL-@bbkk`cjlZL!H{%jh;pG!md^J&Pd;`MyLE^T@L$n26nNPl2tc5#_~aAY=CW*-`v zO)awzkId@#{K&|xe$S7N%+4v_^J62k`ZJ7=%%+#e-7+$pS!N#}neAF;pBR~4US^*h znawViH z{55F^zc~%zuT4Yv>(UVZ`ZR>UAr0YgOhfpa(h&aUG=#q;4dHK1L-^a$5bpCaMJ00@ z)@F0l5dKnsE#97=6XNH6M;gN4nTGIpr6K&?X$XH$8p7Y3hVb{LA^iPm2>(DD!ataX z@DHUS{KIJo|416bKbnT{kEJ1eJPqNuq#^v{X$b#B8p1!BhVWa{5dNt&gnv2>;h#xE z_-E4){<$=Se?AT2Ur0mv7t;{_r8I)gp z!oQV<@NcIf{5xp~|85$>zn6ya@24UB2WbfZVH(1Jl!kEEf~{5Ja{w`${47j<7G7?j zg(-t~{|ooUbMjf3o7*A0?u)#h5xDaKgx{Qo@Onnz^^Cyl8G+X`0$QN_YXPtOBCq=*ulpje`y#LVBCq=*ulpje`y#LVBCq=*ulpje`y#LV zBCq=*AJ6Xz!f#1Kc-k=K2Z*L{)KeUaCFk=K2Z*L{)KeUaCF zk=K2Z-%PeAzR2so$m_ny>%PeAzR2so$m_ny>%PeAzR2so z$m_ny?Tce7nbWW~o12F4$!o#H8QJUcQo9z^_riF7=MX>ZEolg!ye>?fm%Sb@#p}Y9 z!EgJ8*9HIn>fGE8;q|(}>luO9GXlRkkAv`fM&R{~!0Q=-*E0gIX9Ql)2)v#Vcs(QV zdPd;&jKJ#|f!8wvuV(~auM51M5qLc#@Onnz^^Cyl8G+X`0(!)M(}#P)UE}$z&!2|y1=0|{U>d>~N<;X!dFT|_{wPrUnLFUtEM4*wKRmUo`&!> z(h$C88p78~L-^Wh2p>yBCjXmOFZ^e@ygu9Jb%F5h(h$CV8p3x-L->wq2;V6U;X9`x ze3vwY@0y12-O>=gdm6&`NJIFZX$ZIf_cH*W5{K};(-6K-8p8KYL->Aa2;V;q;RmE4 z{J=DXAC!jhgVPXxNE*TqO+)x$X$U_&4dF+mA^gZRgddfL@T1cZeoPv|k4;1PacKxY zJ`Ld~q#^voG=!g&hVYZq5PnJ;!cR>@_-SbfKRpfMXQUzg%ru0bm4@)CX$U_%4dLgc zA^hAlgilLD_<4F@Lq9)ji1`I+2%nyY@C(xreo-32FHS@Fj5LH_l7{e^X$ZeG4dIui zA^h?*gwIMt`0O-Zk-4gnlme9|(gnq6i^z$vDUuX&aVoT_kT0+0v68e>vxTyYiOI%RD z+7kM;me8-agnpwX^qVcA-)af{c1!4YT0+0u68gQC(C@c|{-7l;tiRI|u79^B^!Hlg z{QCPX;onj9gO<=AwuJs+OU(T{tbTtsKhfnGrO&O!T0$@05_*Z2&`Y+2UaBSZ(k-Ev zX$ie-OX%fVLNDJEdWDwIE4GAQsU`HvEumLw3B77d=+#<6uig@Rjh4`BwuD}*CG^@Y zq1R~%y>3hB^;$x&-x7L*me3ovgx;tn^u{frH)#p|$(GQcY6-nrOX$s8LT}L$ddrs3 zTeXDVx+U~BEupt<3B6rQ=BU?fr)e`#Xme9wvgg&+<^l>erk8cTmLQCiqTSA}I68hwp(5JM7KD8zEX)U2o zZwY-yOXxFOLZ8(V`s|j_=d^@AwEuqhE34K9J=nGpyU(^!%;+D{tw1i&N68cg* zQ?vc)ZWFFw))M;ime5zUgub#R^i?gPuWku_O-uZf`r4Lo{koRW*SCbep(XTTXuN6pS#?=UU6P8nAbk1=a~CED_8`S;$lxR@~a@&4!j+RsmZop&|!?ppuexW!sR z`)ADkbF|<8+{d7otO@PDoV!2u(lw!%X$kE1CG;9Cq1S8)y;e);wOc~3(-L~!meA|9gkHZT^ad@VH*5*L zQA_BJTS6~riQlg`X$jXioxM5p=xnrr&>a9))IR2me5LpU@Ke#Fo$}wS+#oCG;sRp-*iIeOgQC(_2EH(GvR1 zme6Olgg(0^^f@h|&us~PUQ6inTS8yZ68gfH&=<9YzPKgyB`u*BwS>O3CG@9TLSNPr z`tp|0SG0t_vL*CYEupV&34KjV=xbX-U)K`)`j*f)w1mE~CG<@#p+D0S`sS9$|F9*Ns6Wx)Us=V(y%XK`@6o!iCG=0Vg#PK4&_B}>`m-&ef3_v`&$Wd9i!Gu5 zQcLKcZwdXETSEVpme9Y@68aZgLjO`rX!pfJ6ZcMZ+qL^zLVw)vi@(~(N%(kwxh3?k zw1oa^Eunw4CG@Ygg#PuG(7(|V`meWy{u?c!|7J_*zts}@Z?}a0J1wFAZcFID*An{v zme3Ehg#P<2p?|X_^l!C<{#;Awf6x;8AGU=4M=hcMaZBib(h~ZgwuJs?EunwACG_vK zg#O)@(7)Fb`uAHx|MQm6pKl5M2Q8uhMN8;^*%JC+wS@k|me7CH68c}ag#I@zq5o}5 z=zrG|`ro&N{tqpoJs0d~!sh|u+Vpc_`nmAu?Q>z_%%lHDUl;wmM(^wGgdSfP^>{|m z&Ig44Y)k0zjG)Ief*#KZdORcO@rhX0^kFSe*d|lMz>!KcC7xnnMsK?huJ-#mL@pVy;uZwzoUDV_2qQ1XBCqh5a5_)`H z)Z^=-9$y#r=lXSo9$y#r_`0aa*F`!KcC7xnnMsK?huJ-#mL@pVy;uZwzoUDV_2q8?ut_4vA| zy)G^=aqmR8UAwO(^z^ylqcgJm#v!&kK4yBk1vrpg-HMBlLJi(Bm0Fk7oouo)PqTM$qFKL62tyJ)RNt zct+6U89|R{1U;S+^mslXb z`>Um-A+6{8TKmRe;KL2B`{nj$x!AI}&=DBy+{=vLocm8ngU(CC7VK65&E-sre=FQx4Ax^+w^#D)8n~K zkLNZ$p4;?zZqwtrO^@d`J)YaN$I`;d@AqCcx$WA0-6r(M{k|B_us+`L46Da8tR7#t z^mvBV;~7?uXIMR+VfA>1)#Djfk7rmto?-QPhSlR4R*z>`J)U9p{rx!;`hk|vdOXAG@eHfSGpruZuzEbh>hTP#$1|)R&#-zt!|L%2 ztH;+dJ)U9pc!t&E8CH*HSUsL$^>~KW;~7?uXIMR+VfA>1)#Djfk7rozxnM^V)6a$J z=Yr1%cx?Q6{9IT#^U=?~Cf+-_?Z0PmU$+T8zAozVjG&zl2>scX(Bm0Fk7oouo)PqT zM$qFKL62tyJ)RNtct+6U89|R{1U)_%=<#(?kFSe*d|lMz>!KcC7xnnMsK?huJ-#mL z@pVy;uZwzoUDV_2q8?ut_4vA|$Ja%De}7JdexN1v_`0aa*F`j*u*F6!}h zQID^SdVF2fhX0^dmfm3F3{t1fgWEM_4vA|$Ja$YzAozV zby1J6i+X%r)Z^=-9$y#r_`0aa*F`@UcG75_$1{Q+ z&j|Xn{W?OAX9PW-5%hRQ(Bm0Fk7oouo)PqTM$qFKL62tyJ)RNtct+6U89|R{1U)`4 z=<$r8$1{Q+&j@-vBk1vrpvN1N}J@KECPa{zvBk_s7TabH5L3hWjULe5{@qJQvNeXmVrv z-_7p-%hr8CxUW`f3B7Vl=v7)mui6rNwU*GUw}f7!CG?stq1S2&y>?6Jby`BN+Y)-c zmeA|Bgx;Vf^oA{=H);vJaZBh0EulAQ3BBp;&6ywnP5sGk5w3r#CG=)3p*L>{y+upt zEn7lw)e?H^meAX@gxTEuoKV34MG^=o4B(pV$)mq?XVpw}d{WCG@E+ zp-*cGeR@miGg?BQ*%JD!me6Ongg&Pw^tmme&ua;NeoN>JT0&pg68fT+&=51+aLOVEuo)l3H^La=oeZ-zt|G` zrIyeyw}gJBC4R2{c1wJ=ezhg^Yb~K)ZwdWIOXxRSLci4#`t6p`@3e$|w*MEup{H5Q!SyNZVCNNOXz1?LO<6M`uUd7 zFSLYyu_g3NEumj-3H?e-OwZ=iv-#4q4~3cMf@Qi*cy3s>CG>JFp_gw7Jw2OG&*t0C zJ{)Eq58HK{@c7uiCG-w0p?7Qvy;Do*om)ch(h_>tme9MkgxG#bqo_#pXJRUCTHsSHHs3rbG zeQ8Vlmww;;-}mePJe|$mf46o&5&E{4(6_gQzN01doh_m7Y6*RJOXz!ALO<9p5HUTQ zPtWGB_BbGXpZT?x(66_IexoJyn=PT=Y6<;zOXzo6LciM*`n{IW@3(~hpe3f?H~;9r zW#Dsr$?g+EFVzy-xqNyypSe3`?vHi4Nw|O3Z3(?zOX&4mLT}I#dc&5`8?}VqxFz(0 zme8BDgx+-a=FE@N@0(B0=F_wJ^lbj0dp7?``#UZFvFEd&(-L~$meBjPgxXZ2s)*!*%9sf# zgnqFl^h+(FUv3HgN=r=tPRsP~v`oKmKKr@gxp;aupPtR9XY-FfACBxM;q&9Dme5DH zgg&Mv^sz0Wk825id`svPT0)=L68fZ;&?mQqKBXn}sV$*TYYBaNOXxFNLZ8_Z`mC1F zXSaktrzP~cEuqhA34MM`=nGmxU)U1*qL$Fk<^QeD<`4GI`?EZo{|~+Y>-u;IAK&#Y zp>JpjePc`fY5#q5@Bh}`FQIR134MD@=sQ|M-`NuSu9nbuw}if@CG>;s0uj@{(=z?O z`Kvt+2!ChfwU*G+ztb}P|AKlfcG#d2f2U=K*@wf-<6*~c z6CNKswS?ZeCG;*Wp?7Twy<1D@-CIKM(Gq&kme6~(gx>DhdGHotuK;qZws^mzEQJe&Wd z{eM9{{-G#c_ojm@1U;DXk6Z-j<&@Z%vez7I=OD&;aZVCNLOH99S{-^VO^TlTG z&Y2$Tdw`xcwk9FVPa(-+3qW`1juR z(!HI~%d~`Ewk7m(EuqJ+5%l;qf*!v{(Bsz#di)weuhRP^^r|hPS8EBqdQ0dvT0*bc z5_+wc&}+AZUZ*AWx-FsCYYDx6OXv++LT}g-dZU)m8@Ghsq$TvGEulZz68cjup*L#@ zy?IOMEm}fv*%ErIme5`uLX6C$xkq0ep!eNId0b6Y~6*An{tme3cpguZb0 z=FE@JO?vKa&(;@ph4Aw)ZV7!!OXx){p)YL-{pps_m$ihxye0G%EupV$34K*d=&NUM z&iweTNzc9QnfjWp5Ptr(EupV#34MJ_=o?x>-`EoRrk2p3X$gIEOXyo#Lf_gF`nHzP zx3`48qb2m6Eursf34M1<=zCg1KiDo1@zkWf4NulzXo)B4FSdmKQcLJBw}k#mOX#n* zg#KDf=&!eg{zgmaZ?=T~R!iuIT0%eE68e#r(2ur+eyk<*<1L||uroE=Pj;Jd{i&AF zPq&1ArX}>VEuo)l3H^La=oeZ-zt|G`rIyeyw}gJBB_6N8-4c)1ueOAKttIs9Eur6N z3H@eE=(k!zzugl0otDt=wuFAKCG`6(p+9Je$LjC2gzMjJ3H`m6c%=S*OFUfvpe6K& zEunwd5)ak>U5+!;tp#$z^FQX=3-hkceLnl8FV4Gm!`wS;Uz&IA{NdWm^R8_=Tzh5S zwVQ@(-=25v!r|Jh^R8VpTzhTawTp*quTQQmHF?{z!?oY++H&)H#d*D8UXLHkoBjF+ zhwE?5`+bb}$y@WTT|4}`xBrA|mrky~^Cw)J`=vkb+PgpI*S$Bn_W8*>k~42=FYU%+ zEuj~03B5#1=p|c1yRYZ&Z@qL)=w(_$`!kvQbJ6}x=Kfst@-?AXXbJ7lZ0^rZuT&Fy z<(AN^w1i%@CG=`7p;vDSy+%vuHCsZj)e?H`meA|8gkHBL^m;9!*KY~EK}+ZjTS9Nt z5_;p7&z2^lw1nQa zCG>VJp|@`dy+cdr9a}>0)Dn8ci zTS6bx68hMd(8sleKE5UN2`!;dYzci*OX!nZLZ8wS`qY-tr?rGWy(RP+EuqhB34K;e z=(AfwpVJch+?LSiwS+#uCG-U?p)YI+eNjv3i(5ip(h_=6OT1rS+7hmRx+U~wEuk-O z34KLN=qp=7U)2)&>Xy*gw1mF4CG>SIp|5WVeM3v=8(Tu()DrqLEun9234KdT=v!Mt z-_{cP_Lk6hw1mF1CG=e_q3>=9eNRj12ipZA7ERjQa7q1zmbke7VoT^RwS@k1OX#n( zg#Kzv=&!Ye{(4L3Z?uH|W=rUAwS<1CCG^8Bp&w}p{b)<*$67)^-V*wWme5bOgnp_e z^wTY&pJ@sGY)k0pT0%eH68eRf&@Z-xeyJt&%PpZ_X^D&KZ@0t+^{XwRUuy~ddQ0dx zT0+0s68f!{&~LYdey1h$yDg#LYYF{+OXv?;;==knE#dliTS9-YCC;zE-xBB5KWGX4 zVN2*Aw#2#h_;r?=``rIva&0khjH_p7=EvXe+Kog1&b;??`Ecz!^RCUke?R5B^R9W@ z7Ji8DO|C69n=?NynB4w;*OndLenq|9yk2o$FPPVJ&k^2y`Q)GV2a}&8XWrDV*bVo= z-2I>zuL-?GOXwwALi-rzJ`TNfP3UD>Li^a}K0fVZoBR0m@-?AXXbJ6OpZoarN;RQZ zZVA0gOXyWwLa){mdi9pjYqW%3vnBLeEuq(L3B68B=yh8{uh$ZK{g%)hw1nQUCGYEupt-3B7$w=p9-@@7NN0rp?7Nuy?aaOJz7HV*%ErMme6~* zgx;qm^u8^j_iG8ge@o~CT0$>u34LHo=!059AKVi9ke1MgwuC;cCG_Dfp^s<@ePm1M zqgp~A-4gnkme9wxgg&k%^zkjBPiP5!VoT_gT0)=P68e;u(5JS9KCLD6=`EqpXbF90 zOX#y&LZ96d`ka=~=eC4CuO;;PEuk-H34LKp=!;rHU)&PuXxV^=oHu&iuIer03qYQhi-l2tWV&me4n}gubyQ zeo)`k60U!yCG^cLp>JskeQQhT+gd{3-V*wbme6;$gubgK^xZ9??`aABV7oxX3X}FW zEMI@2C6=qd*b@3nEup{M68bAGp}*P^`fDwrzupr18!e%~*%JC&EukN33H@+O=to*Y zKiU%dv6j$}w}gJ8CG?Xmp`U6A{d7y{XIer(+Y6? zOSt~smeAj8iKXiAx5Se54_ZQh*b@4OEwMyBeob+{obW%+yEgYZ?ehOL@7kidU)283 z^R8VtT>D?hNq1SH-y+KRp4O>EQ)Dn8*me31YLT}O%dehmPGe1uMJ*z+ZHH-Ua*X|iY z@75A}_m)Ar6u%NTS9-WCG^)@LVu$r^fy~V zf2$?*LoJ~nZVCNJOXx>iLO<3L`tg>~Pqc)7vL*CWEuo)o3H?k<=x19(Ki3la`IgWx zw1j@KCG<-zpb5q_-y@ZOX$~HLciV;`i+**Z?=Sft0nZ?Eur6O3H@$M z==WMezuyx2gO>Q&`a3P*`gdDGf3GEerv83Q{B-?;me3!zg#KYm{8a60mOnBR7EN|~ zHlLo&H|()Mc&^>3CG^HEp%=7--lQe;rn5I^ew?1or)Tqndn6DZABVJrKC~tDVJ)E# zZwY-wOXwq8LLb!<`skL>$Fzh#wk7m&EuoKZ34KCK=o4E)pVSii(-L~$meBjPgxup%;VwsZWA6KC$xkq0ep!eNId0b6Y~6*An{tme3cpgubvP^hGV9FK!9_XL&aNLht{I{yBty-jyw( zuWAW>bxY`LT4Mg+Y56bxzPb1Jr+GGCytDT${k{nA|JIh!x3z@6y(RP=Eursh34K>f z=(}4&-_sKM!FGX&>ECIYe&77n9tVW8@M|rhUvCNhMoZ{7TSC9p68i0y(C@T_ezzs` zdo7{gZwdWDOHBVx%k*r%#_U63=DBUnZWEsC)@lj8c1!4WT0*be5_-Ls(CfE^-k>G) zhAp8tY6-n@OXvkHp*Lv>z3J@DnIETrr)Bzm^MiXN5FQ_gw1hsiCG=q}p$~5feMC#> zBU?fr)e`#Xme9wvgg&+O7CG;gNp%=Bp z5{v!U{hbzP=PPXe@EuI zme9|)gnpqV^ouQ_Uup^ca!cq}T4MV5&8L6geCgST!pw8$GTkOTS1sETdbyU+%eRDH zp(XT+EumLx3B7Vl=v7)mui6rNwU*GUw}f^spPtQU?v9!JW3z4&?w`$DLT}L$ddrs3 zTeXDVx+U~BEupt<3B6rQ=P=YrRo>DhdGHa~Rs;V^S{pMKwb`gd9`oje}?d;kA}dVagHkA?WJ z{rl!ibylCA&8KJc>Dhet^U2wIdN!ZEFMg%_0yoXhCw{f|rpcRUUR!(m_ypm(Y4MiO zOSFVuvL*CVEunqSeENO!nY&}={#d7*g!^aRmeA|9gkHZT^ad@VH*5*LQA_BJTS6~r z3B5^6=uKyD&ipw2J1x`in@_)QKK=XV7f+rOJU{rJ`Ty7Nn=jo9%Vqs~!uz|tCG-_7 zp|5NSeN{{7t6M@}(-J>vf8X5iZ+bSLp3SFc^V!eOkIUJ-PwPMC*2%49W)9K!-ZOWW z_U5_g(PbxZym|6-mz(_D@%_v_hb{jnT$_8oT;Wf+wrKJ`SNs#MeR^_jrFqvjoxI(w zJnx#feZ~*5%H-P2%)aT5|D9Z`P2PCJceZVCM}EulZ#68dLbLjPP#=)c$!`Y*MF{`r>Bf4L>}Uug;b3oW64u_g2`wS?}Y zzh|L|dndZ>+I=mdKkoO%U+v>0e7wKh68cwKLjSdv(7)Oe`qx@Q|9VU4-)IT_*IPpW zjh4`VvnBN3Y6<|C^T3|F$LcziSEo?^{CuhnCQu3wAW&^MG(|`nfRu zT=?_$xv+5N(f^~bi%;~qa9?jH^!U1{$1{RhX0^kFSe*d|lMz>!KcC7xnnMsK?huJ-#mL z@pVy;uZwzoUDV_2q8?ut_5J-h5&D6a(Btc(9$y#r_`0Y+*RLb=_`0aa*F`!KcC7xnnMsK?hu zJ-#mL@pVy;uZwzoUDV_2q8?ut_4vA|$Ja&eb#Z}-dndZ>+I=mdr_Ti+osr!iAGha% z@v(4!f98ab^?{bq)8~bc&dctPkK^;g!kItqAAMeUx{vj~-cIQ8c|nh71U;S+^k@5Z zgdWcbdORcO@r1 zCG;9Cq1S8)y;e);wOc~3(-L~!meA|9gkHZT^ad@VH*5*LQA_BJTS9Nr5_;3wn=?PI z(I32?dzW$RZHltTS9Ns5_;Q~(A%|y-o7RD4lSW~ zYze(nOX!_jLhsTNde@fFyS0Sgy(RP>Eur^p3B6ZK=)GG)@6!@`-agg&?>^dT*w4{Zs3SWD=`TS6bv68gxN&_}g|KDs6JF)g8wZ3%r` zOX%ZULZ8qQ`oxydC$)q=xh3=|Eul|s34K~i=+j$5pV1Qf%$Cq+wS+#qCGme3cqgub{X^d&8!7qx`Gv?cVXTS8yf68iF%&{wpCzOp6sRV|^f zZV7!&OXzD`LSNSs`udj8H?)Mlu_g3PEulZt68h$r(6_XNzO^OvZ7rd1ZwY-zOXxdW zLf_RA`tFv{_q2q5uw5YHsY!bqo~*yn623q3#g@=tY6<=2me5~m3H{ZU&|hl_{q>g6 z-)IT_&6d#LY6<;NOX!DNLO;?H`q7rqkF|t;ye0G#cBW?g$!-&_Kh+ZY>6XyXw1j@P zCG>MGp`ULF{X$FV7h6KV)DrsTme8-X#N+k1Tf+Z_`f5w)*IGiq-V*wame6mugnp|f z^xG|=-)RZ`ZcFI*T0+0y68eLdc&z?TOSt~smeAj8iAU=1x5UHs4_ZQh*b@4OE%8u2 zJv+M(W}iz|Z3)jUtF?q)y(RP-Euq(J3B6WJ=(SryuhSBG-ImbnwS->3CG-X@p*L&^ zy-`c(jax!*(h_>p*_$&zPS4KMv-3ec5(tlvgIhu$(h~a6me7Z_gg(3_^bswgk8BBj zR7>ciTS6bx68hMd(8sleKE5UN2`!;dYzci*OX!nZLZ8wS`qY-tr?rGWy(RP+EuqhB z34K;e=(AfwpVJch+?LSiwS+#uCG-U?p)YKSH5Z$nop0!VAlyGUwuHW^CG=-nLf_mH z`j(c^x3+}7ttIsBEursd34Lcv=(}1%-`x`Wo|e!LwhKf|&(71c^Q%1$2xraLT0&3H z&eOBA$AagMRa!z%&(5>=g&(J9=jqw`pdJT=&yRy!LQl`mAK$aHKi`Y`b0+k~Euk-I ziM1A+e&2a|cAkFU`Kj4Yhnes5KHY7?_kEvf3H@wK=;vBOKi?Aig_h7SwuFAECG^WJ zpGz%gbiVK0-8#eE^PpI}=R)6;p1U^p^{%(B^JA{9JGr)upE%*V$!^%SW#``5 zzSuioZeFiAuNTbg%ZFaC_jgTw_2g|E)DI5V*RRKqVeWeb-n!AeYjeLx*EXJaZTxi$ z=3N_qKby?EHuwATJ~y3rZEtU!@X2}CuARHkern#e`-f}4IrE2e?>&Dwylv6s_U<#k z?xX+Bbk9kwe1Qr+q>%UK4tWme5PKgkGv8 z^wKS%J!a+}H`-%m?(w4imggQXdik2rE3}06Seko0>6L0iuiO%P{8~znUrXuL`gw$2 zy(RP-Euq(J3B6WJ=(SryuhSBG-ImbnwS->3CG-X@p*L&^y-`c(jax!5XbHVZOXyFw zg#J`Z=*?O}Z{8AmiTSA}L68ij> z&=<6XzOW_qMJ=H(ZV7!!OXx){v1@&4OYBmAx+VPl%UVKT-V*wXme5zWgubdJ^wlk) zuW5;$>uXzLr~0~<@bj;434KFL=o?#N$NHw0*kR~vSNrL`_ozS9+X=t^=9bX6w1mF3 zCG>4Ap>J;qeMd{^J6l5E)e`#dmeBXKgnqDHAY!{odmFZ`zt9rf)L(1~{iT-BUv3Hg zm6p(7Z3+Fgme5~s3H^EukN63H?}0=*L?^KhYBU$(GPh zwS<1UCG;~bp`UFD{aj1v=UYO*&=UH^me4P?gnqdt^eZi~b^Yy@*s^}LCG=}8pdR<}S2*|Cnq0{FrO|{+Mg~O|E@@^3VATEpcf#T%P$U zlU}?fTwkIk^pY*1mud;UbW7-ET0$?|5_-9o(95@kUZExQiY=j6Y6-n^OXyWvLa*8q zdbO6&tG9$+qb2m3Euq(H3B7hp=yh5`uiFxOy_V4Hw}jrHCG>_Zp*Ly?y>UzE1ubzv zf6$w>gzKBm-kkYy|E_BU?fr)e`#Xme9wv zgg&+<^l>erk8cTmLQCiqTSA}I68hwp(5JM7KD8zEX)U2oZwY-yOXxFOLZ8(V`s|j_ z=d^@AwEuqhE34K9J=nGpyU(^!%;+D{tw1i&N68cg*Q?vc)ZWFFw))M;ime5zU zgub#R^i?gPuWku_O-tx&TS8yg68id<&^NS%zOg0rO)a56(-QjTme9Afgub;U^ldGn zZ*K{GM@#5CTSDK}68i3z(D$^2ez092V$r0%4VTnkXo-vKFSdmKQcLJBw}k#mOX#n* zg#KDf=&!eg{zgmaZ?=T~R!iuIT0%eE68e#r(2ur+eyk<*<1L||XbJsfOX#OsLO)^{eza!AGU=4 zVN0A_kLP%GgCAnyVt)MSe-pVzPB?JhwM*wNvjzJ+Eiq|x+2r;iGuJ+P zZu{MSZ&wdL@3-q4hyI<(@BQ zwS-=~CGO7CG;gNp%=Bpp^J64Yzfyt-4gn;me7~CgubFB^p!24uWAW>bxY`LT0&pj z68gH9(AT$wzM&=bjV+;XY6<naLf_dE`mUDHcejMTrzP}*?E(?gv-9-q{A!N_!ddgRme8-agnpwX^qVcA-)af{ zc1!4YT0+0u68gQC(C@c|{-7nMXXokJd3tu9{aobZVA0lOXz)DLhsiSdjFQt z2egD<*b@4{me2>agg&?>^dT*w4{Zs3SWD=`TS6bv68gxN&_}g|KDs6JF)g8wZ3#U+ zJOB54cJ}9YMSqTjzOp6sRV|^fZV7!&OXzD`;_$_$-*Dl?!9tVW8 z>GbUU|NedF?xq>8o2)sa=N8|Sp1XFwoN(m4YjdyfE*~}T+J$q!s6Be#wFSeqW9D7E zc(``#ylb}&*N&Ta?ULcz@$;@N8m^r%@0v@OO|G9fx%P*{&!79Xzc#r;E}dNajY-cv zNB^}+-#EGcE0exy(!W}J)8x%F|2=d6xS4-s>%SkqcuVLdT0$?`5_+kY&`Y<3UZy4V zvMr&PYYDx4OXw9^La*2odZm`oE4PGRr6u&LEumLy3B7tt=rvkGuh|lMt(MShw}f7& zCA6)p(3`Y`{$xw&Pql>JtR?j3Eupt)3B6@Y z=&f2pZ`~4lo0ibqwuIiUCG_?!p?7Er?duUjk6)AM@oN(8>k&fxdW6v9*CcxUnnaIZ zlW1R$5Zc!xg!X(j_nf82uSxXyHHjX-Ceh>9B-+;_gkIPZ`oNaZ2epJgxFz% z34K^g=)+q=AJG!}$d=GYwS+#pCG;^Zp^t3|eOycE<6A5Ot^{-29neg_P zwuHW{CG_Pjp|5BOePv7Nt6D-|-4g$~zNRI9yS}z1y!~}8p|5WVeM3w9R()ej{N|wj zZ(Dz_zNz~OKmUs@p>J*peM?K|TU$ck))M;mme6;!gub&S^j$5X?`{cwPdiV=X$E%EF1*IPn=qb2lxEursk3H?Az=m%RuKhzTX;g-;kw1j@NCG=x0p&xGv{X|RX zCtE^4)e`#Yme9|%gnqUq^m8qtpKl5MLQCiuTSC9o68hzq(66+Fezhg^Yb~K)ZwdWI zOZ-~>&6fDp`puTmZ?%MeyCw8HEur6S3H@G6==WPff6x;8!6ZVCNqOZ-axotF6J`nxTmzt2{SZ~-Z!u2n-gx;tn^u{frH)#pIX-nwMT0(E$5_*f4&|9{I z-l`?^)-9p8X$ie;OX%%dLT}#^dWV+KJGO-0sU`HzEunX53B7Ae=-paE@7@x6kCxDT zwuIiRCG_4cq4#MCy>CnC{aQlr-xB(Ome2>bgg&Sx^uaBm4`~T~XiMnBT0$S*68eai z&_}j}KB^`3(Ji5mX$gI7OX%ZTLLc7}`h=FyC$@w>sU`HuEul|o34Lly=+jz4pWYJs zjF!-6wuC;bCG^=Xq0eaveQrzW^IAfm-xB(Qme3crgubXH{#(!H_xAoS@BI+o{}nBv zuWX5X>Z@A9^{ZRrcj{|eVmzCVXY(g|91zZ`Pqu`9swMQ(Euo)j3H@wK=;vBOKi?Ai zg_h7SwuFAECG^WJp1CG;9Cq1S8) zy;e);wOc~3(-L~!meA|9gkHZT^t_hP8?=O;KX-HT<9Ie7&*lg8NFY2u#*GK;3*X-o`hk|v54MDUs3r8nEukN23H@kG=*L<@ zk7x7ob@Ok}JsigGY3bvz*z{xJxzo9PJe$vbKzuyLv-x;7AJ67L{QGcNHwk|~4sQv4 zL`&!+TS6bz68h+t(8sidKDH(FaV?>bZwY-uOXw3@LZ8$U`s9|-r?iAVwI%dvEul|u z34KOO=rdbFpVboj?3U2yw1hslCG>eMq0es#?Ogs}dp19@v+2dXJ@McB`_0GK&Bxcx z$Fuo(Hvd_l&HwG()7|99+4uL2-)}yCPs_SJA_$*L*J}yAeoN?iEulAP3B6%U=r6Q{ z-l!$?#x0>YX$ie)OX$s7LT}y@dW)9OTegJWswMQ+Eupt*3B7Gg=rS^Vey{pJ4tT-DzrLSNky<5_$>i;rjV@htwc zK8t^C?&)Uo<9HUI`(Su3UbdTrUalqd@-3lPXbHVyOX!taLa*EsdX<*YtG0w*ttIs8 zEuq(F3B6`Z=(SoxuiX-QotDr)?X&p&xpzJJas2-B@oauTj|9TwZT$Z7@pbY)>if(8 zThHbmUzhhdBlHz5p|5NS{j)rqAJ+T7p|GO;Xf0yN%9tVW)k$JWy^m8qtpKl5M zLQCiuTSC9o68hzq(66+Fezhg^Yb~K)ZwdWION{^9&+&izx%k{eVe(wMM7IggtxL9q zUaBSZ(k-EvX$d`^&40SDlaFWf@pbe4=l*(4{(g;T^YMFFE*Lx>oXh{Hv-u+9_pprr z&BPub_e~!!&Vu)Mo6rxmgnqCk^g}J7A8rZ#NK5EPTS7n95_){yd_0?vuba=El^5&p z1mW!LTt1%7Cm)W<$76imeEgo4UFRMSlgGnu-6s5f*}Wz79xb8wYze(rOX$5@LhsWO zdf%4N`?Z92E+_POHXqOC|FviHpY(O}lNS2f{`bts?`avor)4~wk7x6r_1U~nlRx0b z>8+D{ZX5o8E~#t$5T|^`wb}Q=ocbBpX7BT~&$xE|+`IhT>7Q}!%IUQ;KI7V@(`#pb z#1*D|&XH>l@aD{z6OWjaou)+!A_|me8BFgx;(r^yV$0w`d8yWlQL-T0(E#5_+4K z(7vyq(7vyq(A&3!-k~M*jxC{gY6-n_OXyu%Li-wM_H~dRejTKTUkB;o*Fk#tb&wu@ z9i)d}2kGJ0L3-bQeM0Zo5_TSA}L68ij>&=<6XzOW_qMJ=H(ZV7!!OXy2mLSNPr`tp|0 zSG0t_vL*CYEupV&34KjV=xbX-U)K`)`j*f)w1mE~B`#d(NB^5Qu79!n34L=*=v!Js z-`W!TwwBPhw}if!5eRoUfd)j#-o}0G&@N9i=OFUD5y(RQFT0-B~68ip@ z&=0hPey}CiLO<3L`tg>~Pqc)7vL*CWEuo)o3H^+n)LcK?ZNl~E zT0%eH68eRf&@Z-xeyJt&%PpZ_X$k#mOX$~HLciV;`i+)&y8dQMJXycl68f!{&~LYd zey1h$yDg#LYYF{+OXv?;LVws2`lFW6AGd`5q$QrJzts}1f4e30r!DbB{hgM0y#8)W z=OUbo`K)XU(@-LciS-dOSOS>Fz!Z6K2nY;^LkQovW8NvEj#{~g;a2fg~gd2-jQrhQ&Nr|bXjifLanc^_Z;PnXyC4Bp?Rwcq>V)BA?M z_cgmV{Jrbp?_JNE{yd|HzxS(azxS)AZ*y&Z&-AvrruJ*kKBmKe&g&Mq_7!hB<<-eg z_~FmHsQYKX?)3}&JU?abpS{;X7Pz+i^!=F|7r5r*H{+%Su3bO7(0*}&YqMX=o7}v> zHNVdpw@j}sHoar`ao@0O!;iZje%$rE**j|W{6XJV-!tf2>*2@w%E7hUr`J}P-uESM zIlaAU*M6zng#L0%=wEIL{VOe@f3+p_pJ)mFYb~LFy(RRYYzh4vEusHZOX%Nh3H@6w zp?|w2^zXEU_Ho+R#FwVJ?b-)#x~=UPJl`IgYX*An_K zw1obPEusHXOX$Dc68f*Sg#N27q5oP-=)c|)`fs#^{z^;ezu6M{_gg~$t(MS#yCw7= zw1oaUEusHzOX$DX68i7Ag#HIDq5ok^=zr7_`X9H1{wFP=|7lC;f7TNEpSOhm7cHT` z+7kL-wuJsyEusH)OXz>o68hh^g#KDf=zrG|`VU(||NEBE|Dh%He{2c;pISnDF4)$D zzXybC&j@;WM$p4Ef*w8>=;3pL9)4cb!_SL)_<2zeKQHRx=S4mI zyr_qt7xnP-q8@%;)Wgq&o)Pr$jG%{S1U)<>=;0Yb56=jCct+5}GlCwT5%lnkpoeD! zJ$zo!!!v>&o)Pr$jG%{S1U)<>=;0Yb56=jCct+5}GlCwT5%lnkpoeD!Jv<}muk`0c z=;3pL{{8MJ^ze+Jhi3%+2fZDkhi3#mJR|7g89@)v2zq!%(8Du=9-a~O@Qk2`X9PVw zBk187K@ZOedU!_A!!v>&o)Pr$jG%{S1U)<>=;0Yb56=jCct+5}GlCwT5%lnkpoeD! zJv<|5pBMHq@ujJ5yEZ%{=r4Cap~vTfAHI&}5ZNHxZvWs`P{D4zIxg_E;V=KhH39Ky?*1gf3NF)&DpQ(eS-2)#U!*1UqAj5pYYDx0OXww9LND18+HZOK+n)5&HR1X)Euoif3B6oP=;d2N zuh0^D#g@=3wS-=|CG;vSp;v7Qy;@7?)muWZ(Gq&ime6apgkHNP^g1n}*KG;CUQ6ip zTSCuk3B5r}==pOuCqLd%Z`duu^)Ix9-l!$?#x0>YX$ie)OX$s7LT}y@dW)9OTegJW zswMQ+c2aY_O}7cxw`~c%T}$ZgTSD*95_-p$&^xt+-nk|8E-j&VZ3(?wOX%HOLhsQM zde4^7d$ok#yCt;q#OzF=_pJ%NUrXryTS6bu68gZF&LpU@Ke#Fo$}wS+#oCG;sRp-*iIeOgQC z(_2EH(GvR1me6Olgg(0^^f@h|&us~PUQ6inTS8yZ68gfH&=JcwuFAJCG_(xpu`ui=hQ9Yi`eGKNFOO|N~&n?TggkG*C^ztpCS7-^n zVoT_iT0*bf5_*-E(5tqDUackc>MfzyXbHV$OX#&)La*HtdYzWg>$ZemuO;;QEurVN zgx;Vf^!&M-lOM;k`FJ)zphp7X@o`{F=!059AKVi9ke1MgwuC;cCG_Dfp^s<@ePm1M zqgp~A-4gnkme9wxgg&k%^zkjBPiP5!VoT_gT0)=P68e;u(5JS9KCLD6=`EqpXbF90 zOX#y&LZ96d`ka=~=eC4CuO;;PEuk-H34LKp=!;t7Cq0|r)z81IpGSCqm$!tzq9yc| zEupV!34L`-{Bb{<-`Vf?mVTdv-~X*G@$Toxv-!h)90+I2M_NKZ+7kM)me7y4gnptW z^ph>2pK1yHbW7-GT0%eD68gE8(9gGoexW7wi!GsFY6<;vOXycxLciJ)`n8tOueXGL zqb0_(`FJ*8eD0wzdG1-F+l1$$C0jx-)e?H?me9+zgkH8K^l~ksmv0HZLQCisTSBkY z5_;v9(5tkBUbQ9kYAvBxZwb9dOXxLQLa)^ldhM3b>$HSkwVJp|@`dy+cdr9a}>0)Dn8`K)uS>tx5_){yeEfcMj|I;i2EunX93B6lO=-pdF@6i%^ z&z8`8wS?ZgCGXR*@pK1yH zbW7;*Y(9R!`S<4@4nO$swDhp}0T)eg{p$42+4t++H2G$bAO8Pjd9&HITNb!>>)_h2 zP2SI!{7-Ry!0hM!`t;6grdMuT;QgHUgP+lVb9wrHK0m#0_rbN__%YZ1)Q`FLn?L5- zZ!K`mFE?-cYyS51+P|E(-_z{(b#XUbp1et`7j6l?NK5ENTS70^657`vv#&w)@OyW( z@7p1?598!RIql)UA8>v6y*qmNy*qmNy*t|X?GSo}me4D2v z)*E(3B5y0=p9=^@6-}{=a$gBw1nQZCG>7Bp*=Uuo*(ocHKF%x3B6ZK z=)GG)@6!@`-ciTS6bx68hMd(8sleKE5UN2`!;dYzci*OX!nZLZ8wS`qY-tr?rGWy(RP+EuqhB z34K;e=(AfwpVJch+?LSiwS+#uCG-U?p)YI+eNjv3i|wT5`jT!Fu3y>`e`cW{{l0qF zFYkUrU(pi!%9hYqwZxz9{axJ>u3s~EbMoWD)1KXTNqucs2ycH~OX%xcLf_C5ztgXG zV@tSxQ%mSCwuHX9CG;&Vp>J&oeOpWD+gn24(GvR3me6;#guc5a^gZo75f@F{eYmi` zw{DE_jeb#cCkCB z{J90Lx%!p$}^beRxagBU(Zq*%JDwme5DHgg&Mv^sz0Wk825i zd`svPT0)=L68fZ;&?mQqKBXn}sV$*TYYBaNOXxFNLZ8_Z`mC1FXSaktrzP~cEuqhA z34MM`=nGonFD&$<=X%#K>V86B+!Fedme7~BgubjL^yMv~uV@Kw<()GP9u~hx_meAj534LEn==)njKhP5T!IsbuwS<1SCG;aLp&xAt z{a8!r$6G=_(GvQ}me5bNgnqgu^fN7?pKS^KTubQZTSC9k68goK&@Z)wez_&|D=nd4 zZ3+EaOX$~ILch@xOV;0PiN)(TTSC9p68i0y(C@T_ezzs`do7{gZwdWDOXv?4R?{`cl`WyK zY6*RHOZ>%!#dXG`e4T0-yL5_+GO z(EGN8-mfL}{w<*oXbF8_OX!1ILLb}``jD2;hqi=1tR?i}EuoKS34LTs=%ZRfAKemq zJUjp5o}K;qU(@4&(ATzvzOE(o^(~=qXbF8|OX!UzV3W~9|ywgo)5Hyey}C< zLoJ~nZVCNJOXx>iLO<3L`tg>~Pqc)7vL*CWEuo)o3H?k<=x19(Ki3la`IgWxw1j@K zCG<-zp>#KCBbYo38oG{oVJP^z7ONa>`#> z;M%;|h4!y5aBcS7<4ykB0@vmbe%@bS;M(jt$(#I*1+LA`6Rv$_fotcubIRXb;M(x_ z^83?ki%#Em$#!1;?&s(2`M0~a)B?Ti0_|-%d;5C^{kNv?`)l1d2svCZ}$uHUJgy=hD6&00cl-V%C?me5zAd5mYYDx7OXvey zLLb-?`kz8yt;q5PN34K{h+*@DX60Tp-68g%P&{wsDzPcrT zr@p2oe!IT5CA|H0EupV(34KFL{8oKqOZ?`buUYArcmKWmrtT;F{4chIzPTmzEiIvM zZ3%r_OX%BMLf_F6`p%ZnceRARyCw8J?K}~`F>UwZPuBOg#IM(1ZwdX4meBXLgucHe z^aCxSA8ZNzP)q2CTS7n568h1W(2uo*e!M006D^^iYzh5TOX#OtLO;_I`q`Gy&$Wbp zz9sYvEumj*3H?$_=$BhUztR%=)t1n&wS<1XCG;CD@oV)rTjE#iH(Nr#)e`#cmeB9C zgnqXr^m{F#-){;1K}+ZlTS9--68hto(4Vx#pQyjp60U!{CG@8)@hkOrTH=@M@3w^g zUQ6ijx5SrgU)OxiOu1;fLpU@Ke#Fo$}wS+#oCG;sRp-*iIeOgQC(_2EH(GvR1me6Olgg(0^^f@h|&us~P zUQ6inTS8yZ68gfH&^{Nk-F$rAe7znDgy-J%TSCuk3B5r}==pOuCqIs7 z^YLtcK#v5%*jv` z@NBMsnrHK$VEuo)l3H^La=oeZ-zt|G`rIyeyw}gJB zCG@K;p$k^O^k_J>LzNXYWTZToZbc zme7l~gkG#A^zdsOy+rpD+SfOP_Vo>+hhO99;nz5N_%)99^$nqiU*qT%x}VU)uW|J7 zYaH$C8$z$r5_ODCG-s~p>J#neN#*5FSdlfxh3>1Eun91 z34L2j=-XRD-_a8K&X&-3wS>OACG{f68h7Yc%uGJOFUkGwnU)>V=nwHSlwuHW}CG_6XyX z*h$Uxv)v|Kf379;^DUuYXbJsdOX!zcLciP+`jwW@ueOAKttIs9Eur6NiSg|G(>Xi$ zVVy90o)G`G=L+ZQ*|phokbA$jz_qWqbIKYET)Sy@q5XHCac%Z~|L`-e&71z*fBzZR zX7Bbt{FrP1ae-^IU(3(^PYYc8%IqDsYc6oj=Y!dE`hQ;F+Kq#s_dhLg?UJs|e$9Ip zxaOzL-sip3Yl}_a*OJpU_jPTl1$x;9dfuSF(c8}-^!@ccgZ_Fw{Bs@tnLn_=wc(%L zg9}`n{d0M54=r$QcXv*Cc!6tQnO$f%oL-wf|G4C1{pcf){D=Q-a&7K&ul`@B?bp@+ z%d~wi)c^CeZ3hB^;$x&-x7LWOXv;eZcZK#-<|gCKG!$w zI-$QXcXRUg-o_1v-@1%cyRsO)1KYu`X+W}ozVNWgx zTt9bk{mp65?sNUTt`qwFme3c>-JJaT-<C-ikKp|5WVeM3v=8(Tu()DrrOEun9oyE*y&zC7*OeXifqbwb}dcXRUg zFHL)PpX;}EozS<>-JHDri_@On=lUIk>n}`ucAx8a4z52x?b&^<-!-`Y+_Y!+xqf%o z34KpHPsB6Rb|0Ru?`?^v>aVwi{zgma`&vTZ-xB(Pme3EjgnpM9OY>CI~H(Nr#)e`#cmeB9CgnqXr^m{F#-){;1K}+ZlTS9-- z68hto(4Vx#6ZN-R!u4;rg#NT89;?685|7s3Z3+FomeAjCiAU<;*Ll0kDgXNd*W7*M z^xFS1z4obxUktU(*u*XMJr;c>C*GLSNq!`i7R+w!X0?wi&elJ*C5441SuqE_EEukN73H?Y*=to;ZKh_fZ@s`j}w1j@LCG=A* zp`UIE{Y*>fXInx)*An{qme4P>gnqFl^h+(FUv3HgN=xWhTSC9q68iO)&~LQF=JhvQ zV$=G~me6mtgnqjv^gAt~-)#x~UQ6irTS9-(68giI&>yvg{hH9~#`SkwLVvF%^!Hn0quQDNYi7zt(;d&|%=^wf(*bzyIr6 zLSNq!|95>uON?jp@ofG~j|0M4cRZVqXY=uFKKIMfzyXbHV$OX#&)La*HtdYzWgwwBPh zw}if!5eRoUfd)j#-#@Efq*UjJTaX@&T`K^}FZ?}YgrzP~eEuqKP&F4NA zO9B^LaimY?)>bMNn`c{cwqy}z6LeG%UOEiIwPv-$Wv zEicVI944>JzuaxY>-4X*gnqRp^lL4lUvCNhMoWy}(=vWf%i?nng~@Zz65S>|7cJQm zda0JsOSgm`UpF6LH{WvZ;V^kTY}IYT<6~<(skz>!+l1@ewuIiUCG_?!p?7Ery<dXG`e4T0-yL5_+GO(EGN8-mfL}{w<-$v-x;7 z|F1ope{SLbw(n_icK#pvZ0_&Vcs3uu-~5HSr`O3@`NeJ%&d$!|pme8BEgxJcwuFAJCG_(xp`FXev-#x1F*#d*&`rWw`@@#dAGO4IHXr{x zEi27E6egdGR_-?8bBA;Jcs8GWI3^#D@oYZ6Zhp|*!(sAxIJnz{zb}Whgg&$-^kFTb z4{r&5L`&!+TS6bz68h+t(8sidKDH(FaV?>bZwY-uOXw3@LZ8$U`s9|-r?iAVwI%dv zEul|u34KOO=rdbFJC~1V^T~%}^6|L3o5UZt?`c`6v--{bz6gJQx3q-5wI%dzEun93 z34KRP=sR0N-_;WO?v~K^wDUxa-)}yi&EM>CKsXD()e_qGniKk+meB9Egnq9j^!qKL z$Fup|$71|%D31SI(0M%+2+z$Mw1l2NcXRUN`2FVN_p}_)BZ2VvIItzObNP5SpL{qb zACK{D{zw0JS{7dTzwduT@#p=zxxXJ@?4d#E@oYZ6ZvMjD!(s9|`itErysrLIOX!zd zLch`y`qh@uueF4Jy(RP;Eirz-`A_Hj&HJ$afD5L#7VWtXv%mXaSxWuz{{v0UUPJhs z3tXG;;*=#P@8?T*4}Xn|r`P`0{4&g60XmmyE*ys%zDFa z5w3rsCGxO0))M;gme5DEgg&w*^ieILk8TNlOiSovTS6b# z68iX-(Ei;3q5Zo7LZ8$U`s9|-r?iAVwI%dvEul|u34KOO=rdbFpVboj?3U2yw1hsl zCG>eMq0es#eL+j;3tK{8)DkN$+Lf_NQ6LHbB-G>Y7dt2gy`s*#BztIx+Wp&x1q{cua@M_NKZ z+7kM)me7y4gnptW^ph>2pK1yHbW7-GT0%eD68gE8(9gGoexW7wi!GsFY6<;vOXycx zLciJ)`n8tOueXGLqb1I-zu6M!)^E0ieyb()+byBrX$k#qOX&AnLciY<`h%9xAGU=4 zs3r8rEulYYiSz1jwS?>6ZVCNqOPo`GrzOs=zuOY}do7{A-x6olzK-~+2|2l|^>=Hb zme32g{D0ZI3pdN}!@l>nm~q84VrFKoR?N(>tJR8`S&kh}n}$=a8z@5>r)|L-X#NB4`!#b;{PWzco_n7pbz7hJ^St+S-t(QAojrT@%x^{$2YQ*7(95=jUalqd z@-3lPXbHVyOX!taLa*EsdX<*YtG0w*ttIs8Euq(F3B6`Z=(SoxuiX-QPD|)@T0*be z5_-Ls(CfE^-k>G)hAp8tY6-n@OXy8nLT}mKWYj6aZBh=T0(!? z68f{2m|QzguAL{>&NKfOJh$)AO+xS35_+eW&^xz;-lZk!p$}^bJ+~$F;Vq%( zwS+#RCG?Rkp^s__eRNCcV_HHV+YYM?zoS z68f5!(ATzvzOE(o^)0dX(v$ByFYbPra%aTk+Ie#A{C1B6!nLXEW&!a4a&WfptTRXB=d1OWaCYnHZPy>2eQnf!-|_FgYB!du z3GLOV-(=Lww1o3zTS70_5_v+9kP3VnULT}s>dXtvWo3@1B ztR?hRORQdR-V)BYn7KLq>;CnY-6EWC)e?H^meAX@gxme3cqgub{X^!%343tB>7(h_=MOXy2mLN96weOXKB%UeQU(GvQ~ zme5zVguc2Z^ffJ^uWbo^T}$ZeTVjKyfBZV$`Hj7t&^NV&zPTmzEiIvMZ3%r_OX%BM zLf_F6`p%ZnceRARyCw8JEursi34LEn==)njFK(wqtTJjZ!^-v7TVkd98!e%~*%JDJ zme3EjgnpTR&`Y<3UZy4VvMr&PYYDx4OXw9^La*2odZm`o zE4PGRr6u&LcDiPIwQdv6S8oZuMoZ{5TSBkZ5_;{H&~sWsuhSBG-ImbnwS->3CG-X@ zp*L&^y-`c(jax!*(h_>pme8BEgq~`N$+h$3+Ien|1H$9u@Rrc?T0$Su68gxN&_}g| zKDs6JF)g8wZ3%r`OX%ZULZ8qQ`oxydC$)q=xh3=|Eul|s34K~i=+j$5pV1Qf%$Cq+ zwS+#qCGme3cqgub{X^!%343tB>7(h_=MOXy2mLN96w zeOXKB%UeQU(GvQ~me5zVguc2Z^ffJ^uWbo^T}$ZeTSDK^68grL&^NV&zPTmzEiIvM zZHWz+o?JUW*!@7bc6+EL^usNoA885wXiMnFT0%eG68edj&`-96eySz((=DN&X$k#o zOX%lXLOhNq1SH-y+KRp z4O>EQ)Dn8*me8BDgx<6z^kyxgr&?li?L4`5p4;Pq@c5WqJ5R2iC)dt1{~k8#TKMK} z68?SM(h~aCme9AggucBc^c^js?`#QuS4-%-TSDK{68hej(D${3zP}~(;&w{Jhbln7dxMtbw2!fx9oi{9-MD6>*E;yxwo2iHhy08=3CD?`|99*{y`su`*Qs8 ztMU8YrnejZdA#oM&!dNb9zEq9jL?h6(f)qz_xjVp%f{ae`t8wQ=4GR^-yQXJqyEiN zk3T>B#;C`Chrd;O(dflr_?2$>?WTXrR=XRf?}||`(-O{?Z3*o@8Q(8@`I^uxw1jqF zjqfk*z8c?OdgYqXzP1{FeWl%(nlC{`brPKzS6$7BJ`Y=(8I5< zw6CoQJ^cDg``U`o!>_OO@aro*{Q61{zrNDLudnp*>nlC{`btmr`x1Kdme51wk@81&ofR@k)wuC;YCG^29p$};ZeP~PQ!&*YmZ3%sNOXztm zp^s<@ePm1Mqgp~A-4gnkme9wxgg&k%^zkjBPiP5!VoT_gT0)=P68e;u(5JS9KCLD6 z=`EqpXbF90OX#y&LZ96d`ka=~=eC4CuO;;PEuk-H34LKp=!;rHU)&OUeoN>DEuk-I z3B9l-^rbDK7q!H%*O#@#uho~ggx9~KCG?dop|5HQeRWIdYg$5I+Y1tmeN#((aj75wdg^a?eoHSW{??$cU+XIe{7!vqFDJae+gd{3-V*wbme6;$ zgubgK^xZ9??`a8rZ%gR=T0-C75_)kvCE}My?Pd6-`s*$6i}g2JLVvR*^aCxSA8ZNz zP)q2CTS7n568h1W(2uo*e!M006D^^iYzh5TOX#OtLO;_I`q`Gy&$Wbpz9sYvEumj* z3H?$_=$BhUztR%=)t1n&wS<1XCG;CDq2Fu?{Z>o-LjA3l___M+meB9CgnqXr^m{F# z-){;1K}+ZlTS9--68hto(4Vw~{$im7pe6K%EulAR3B7Sk=uKKeZ`u-ivzE|PEulAW3B6@Y=&f2pZ`~4lo0ibq zwuIiUCG_?!p?7Ery<dXG`e4T0-yL5_+GO z(EGN8-mfL}{w<*oXbF8_OX!1ILLb}``jD2;hqi=1tR?i^me7Z{gr3(D`iPd$N4A7M zswMQ%EuoKT34Lr!=;K;KAKw!CgqF}JwuC;ZCG^QHp-*WEeQHbS(^^8G-V*wZme6Om zgg&b!^w}+;&uIyLZcFI%T0)=S68eIc&=Tz*W%ChI3QfxKHn1h zg_h7SwuFAECG^WJpG)hAp8tY6-n@OXy8nLT}m7*b@4pme3csgr46LdO=I*|C!g~f6?!E zWxp@s_y1|S7XS19ym$2HCj9yDY>8v*yIR8e-7TT-X$gIAOX&MrLf_vKdT~1?V)A|Q z$@j(I?r}i)KJq&)q2Fx@{a#Dx_gg}L&=UH?me3!yg#NfC^d~K$KWz#9Sxfw+zAxT~ zG0pgEcCqc!e*MwU9jQzG65B1|?3S6MFKj>SZ2Wx%$2-hA8^4XS9cP_++4y{?C7fO1 z6=LVlIosuP&UT%3cKyt|oBnmTS!b6G&UT-5w#DFVk6C9624{QDI$Jb2+iP_8$D^0| z`xyV7yx5JUYC?PU={FhmGA-eJ*_P1DwS@Nd%lK;;?dumpuhEQ z)Dn8*me8BDgx<6z^kyxgH*X2;F*SZ%=`CwQZ`Bfd>z2^lw1nQaCG>VJp|@`dy+cdr z9a}>0)Dn8h{b6P^5 z+YbxY`LT0&pj68gH9(AT$wzM&=bjV+;XY6*RFOXyo#Lf_gF`nHzPx3`48 zqb2m6Eursf34M1<=zCg1-`f)UzLwDUw}f8YPKkJ7)Lw??>#w)ObM-e`LVvR*^aCxS zA8ZNzP)q2CTS7n568h1W(2uo*e!M006D^^iYzh5TOX#OtLO;_I`q`Gy&)Mmk>F2vm zIDerf^ouQ_Uup^ca!cq}T0+0t68g24(66_IexoJyn=PT=YKdp-Z?(kJ_1i6>-)RZ` zZcFI*T0+0y68eLd&>yyh{-`DN$1R~hX$k#lOX$y9;+gu}E#dq-Eup{L5>M6NYl$c8 z@3(~hK}+Z#w!{;)uS>pWLQcP^^>1sbme5PLgkGj4^s+6Xmum^Vd`svRT0*bb5_+YU z&?~ouUZo}Usx6^cYYDx2OXxLPLa*5pdaahwYqx}+(-L}}meA|AgkG;D^!hEKH)sjH zVN2+ZT0(E!5_*%C(3`e|-mE3`<}IPOn7KLq>zw}L_4sAZx2y@hRZHltTS9Ns5_;Q~ z(A%|y-o7RD4lSW~Yze(nOX!_jLhsTNde@fFyS0Sgy(RP>Eur^p3B6ZK=)GG)@6!@` z-O7CG`B3&XZf^;FM@#5CTSDK}68i3z(D$^2zPBaxeJ!EyZwbA)of0v*cAi{2 zzun`2aIN`HOXzo7LciA%`u&#BAGC!2uqE_IEulYd3H?b+=ucZhf7TK|scYvxjA_Qt zdt#rS`&?I#&u)<;_MLTh?f9VGZwY7P_q+cR&aN1}?Eyc>F%+&N}<*=xxjcOZoLj|F^~~$`Q_1 z8b_;Foz+vbdW%6H-tW73(BJ6o|E7Mg2YcU3`n~>m^!t0fj2|;6^l|x6razokyNd|D zOiSoxTS70_5_2EunX93B6lO=-pdF z@6i%^&z8`8wS?ZgCGagg&?>^dT*w4{Zs3SWD=+ zEujx@2|ceR^bswgk8BBjR7>ciTS6bx68hMd(8sleKE5UN33j?>`owM%&QEFyeR50a zQ(8iw+7kM-me8lSgg&Ds^qDQ8&uR&Mc1!4UT0)=O68gNB(C4>=zMv)ag)N~kY6*RD zOX&G6p%=7-zN97e!j{mNwuD~P68f^1(3iJ_zM>`cl`WyKY6*RHOXzD_LSNew`ns0T z*SCbep(XT!5eRoUfds;%@+Y)=(h~aBme8-Y zgnqpx^cyXq-)sr}R!baHf2$>qs^4x2{Z32hcUwZg*An{ume3!xg#NH4^hYhBKW+*A zNlWNYTS9-<5=YnHZVBh#X$k$^mN>HhUP~NNf4?R44_ZS1uqEczuFJn@;+C12{`JIJ zXNv}BC(SxrFgQDTbhg~+Wh)NO9$o53*S=n?SDke}HLJH6^r^k?#e+U&^n0x_dfAs+ z;>q6k&$NX8a!cs1w1ocIme4=f68h&`LjOWb=wECJ{Yx#Of4L>}ue5~zTP>mgc1!4A zZ3+E%T0;A~Jj}$GM!N0nr|rJ@wf@|MKmV_{g#L||(7)Ld`nOs_|8`60-)RZ`yDg#r zZcFID*An{gw}k!&EusHmOXz>p68ayvg#ITjp}*P^`k%Ih{=Jsazuyx24_ZS1VN2+L z))M-kw}k!|EusHqOXz>q68c}ag#I@zq5o}5=zrG|`ro&N{-c)A|Dh%HAGd`5T1)6Z zX$k!wTSEV*meBvXCG?-Rg#NRZ(Ep_+^nYy${pT&A|Dq-Ie`^W--&;a^F4)b4e+PuK z$-jlkzlEQ+e+!3AKl*?C^Wro8TNr*`)L-uP2tB+;(5?>%{c|m$ht~*tc#WWk*9dxe zji86u2zq#ppoiB8dU%bXht~*tc#WWk&jot;T%d=a7xnP-q8@%;)Wgqa>aX_qMCjq?Mg4odoY2G1i+cEZQU77D zN9f__MLqnysE3~y_3-ne9)4cb!_SL)_<2zeKQHRx=S4mIyr}=EkCV`z2gc6@diZ%! z4?i#J;pas?{Jf}#pBMG;^P(PpUev?Si+cEZQ4c>a>fz@_J^Z|=ho2Yq@bjYfd2!0b zmqxnnZ1P<2qibaM$N%Zi1;fX}S9=T){_KuDZC-m?dK@YDH^dI(mgdSca=;1Yj9$q8p;WdIDUL)w?HG&>qBk18Zf*xKY z=;1Yj9$q8p;WdIDUL)w?HG&>qBk18Zf*xKY=;1Yj9$q8p;WdIDUL)w?HG&>qBk18Z zf*xKY=;1Yj_Icq@6JHwXwzJ_ig8p(ZC-mgG;78v_bASBr{#-EmxBsK-0Dqr94gdDP zIP=f>FFyy2|809-SblV7rbf5#>V8>qblaT1^pF2{4!!EE^Ql?A#i0Fw(BeU#*4z8N z#~+uv#V>Kj63%X#Ir_qxOE|l2bavJf&aN4qojvPp!RY1YoY7f-{d5MoIiTq@t<-0aXP+m*4dTgH)=1Mb$0#W?BZEx*AC9+&pI2v-GW(X z-faA_U$TU=@tpUVd}EaL_l@$HSkwcgg&Gtw6EL7U%TnU zYC_L#34M4==y@%nk7x;fWJ~CyT0$S)68e~y(8sofKCUJ7@hzcGXbF8{OX!nYLZ93c z`jnQ?r?!MXttIs7Euqh734LZu=(AcvpWPDroR-k%wuC;fCG`0%p)Y6&ePK)Ji&{co z+!A_zOXvkHp)Y9(y|5+pr7fWswZvZaWi7F1eR)fG{VQ5RU)d7+s+Q1Kw}if?CG@o| zp|5L+pQ*2JiQVfPTEgq!*b@4tme4o1#BTL1EwSsM{r{=ix4yNP6W;!|me9AigubIC z^qnoC?`jEscT4DdT0-C368gTD(D%25UffQJ*m=}mhMnrKx5SS1H(Ek}vnBKcEukN5 z3H?w@=!aWEKhhHV(U#DUwS<1WCG-<5p`UCC{Zvcnr&~fl(-QjGme9|&gnqsy^b0Ma zUu+5eQcLKUTSC9m68hDa(66ysf{;(zVM=ha0ZVCNKOXyEqLVwm0+t=T23FqHw3H{xc*tY&&OKekrza{h! zT0;M@CAO}8UG`^Y#G=touEi(U;+ynXAUxM@+7f!Rme5lzp*L>{y~WJU>0j@ux9k?- ze5;nwTepPXrX}>YEupt-3B7$w=p9-@@7NN0rnCSuLT@ZV7!(OXzc3LZ8fXInx)*An{qme4P>gnqFl^h+(FUv3HgN=xWhTSC9q68iO) z&~LPaezPU?TP-oUHlJLZFF*58n11eDq1%M#(iK}muhbHH<(AN^w1i%@CG=`7p;vDS zy+%vuHCsZj)e?H`me6xrLa);jdfk@L>$QYlza{hrEulAT3B6HE=#5)KZ_*Na)0WVi zwS=B(3B7qs=q+Y$PX9W&HlJLZ&+CyuczjH*%_rC9|D$X3UtIcst?!#puFWUkH-EC{ zEyA_w3$X&w#ws68hbi(C@W`e!nI32Q8sLYzh5QOX!bVLVwZ{`qP%spS8r~ zXIdsdUq1Q1`OLos&+S`xlhE6=gxc{ofz9**iZ z;qh^FOXy=-LQk&6|IfJ=|DXNy-4QyBoeZb5~41xApM&0^zx? zZ+%QV?q#|`IA69U^l~ksmv0HZLQCisTSBkY5_;v9(5tkBUbQ9kYAvBxZwb9dOXxLQ zLa)^ldhM3bb6P^L(-L~!meA|9gkHZT^ad@VH*5*LQA_BJTSEVJUyE-s^RB0Vo&0?H zTSA}L68ij>&=<6XzOW_q zMJ=H(ZV5fVCG>)p(3iA?Uf2@)(w5MRTH>#IZSMVD)%zjz)h(f~X$k!|d2RlFpBQiM z_aP?#yDXFcU6#qU`Q&F3*18o6kHLd@h34K^g=(#PS4{r%QuO;;4+I;deEdQfxbKk_dtj7uQlm6dKoP6JW@_qBk_su6i z!}6V(r=#ib!+*Ei#N>Z7aq@lh$@k4?{w=t++pe30-o7RD4lSW~Yze(nOX!_jLhoXy zYo>SYHsO4?me9Mmgx;ei^qwuD_i72fcT4DfT0-yJ5_-Ru(EGQ9KA&Huk&oBvAJ@>h1h z5$?bLmA-H8gnqFl^h+(FUv3HgN=xWhTSC9q z68iO)&~LPaezPU?TP-p9`SQP>pD*t}ul*&rjBb5%B{zQj|L5BEqq7H=aCXz^?7=0R zjsH90LrXZjV)S+oFX3$b-!p%;=h@jG%lP+ubO~=ae!hEb*4g;a=Dq#etoONv7moP# z&pG>z&pG?e&pG?8(b?BVzlZyLe7|4TjUL*6`SU*7t53g4tCwjBy=+V9;m`Z%;m`YM zKkGy26G)hAp8tY6-n@OXy8nLT}m;(0jFn-n%9AJ}sg5Z3(?!OX&SuLLbl)`oNaZ2epJgxFz%34K^g=(#PS z4{r%QuO;*mEuoKW34K&c=%ZUgAJY>0*p|@8wS+#tCG-g`p-*fHeNs#4lUqWc(h~aA zme8lQgg(6`^cgLo&uj^OR!iu!TSA}H68hYh(C4*;KEEaO1uda3Yzci)OX!PRLeFmr zy`UxZB`u*Bw#08Q{o_CH<9tysC-h}@x@P+FZWGS0XbF8~OX#awLSNky`kI!|*S3Vd zt|j#KEun8{34LQr=$l$X-`o=VmX^@BwuHW|CG_nrq3>u3eP>JPyIMlu-4gnqmeBXM zgubsO^!+WN7q?R)7LD4=aB2PZmRMMSqb2k=TS7n368gcG&=0kQez+y{BQ2pHZ3+EY zOX$a2LO;f;P=C86oPVb!^mkk0;`)0naZ&yKme4wS-=~CG?z@(Cf5>UbiLmdM%;XZwb9YOXv+-LT}U(dgGSRo3w=9v?cUrEup7c zVsh;~xpto0nCSuLT@ZV7!(OXzc3LZ8o?JWM-~B+ie-^h>A|}_)lWXU2 zEunX93B6lO=-pdF@6i%^&z8`8wS?ZgCGagg&?> z^dT*w4{Zs3SWD=+Eujx@2|ceR^bswgk8BBjR7>ciTS6bx68hMd(8sleKE5UN2`!;d zYzci*OX$h9^H1-!v%mlAdmIq@hL+GbwuHW^CG^cLp>JskeQQhn?$VR*J3rX{K={7r zLoJ~nZVCNJOXx>iLO<3L`tg>~Pqc)7vL*CWEuo)o3H?k<=x19(Ki3la`IgWxw1j@K zCG<-zpXmCk zuhJ6QeLuec^=dVtS8oZuMoZ{5TSBkZ5_;{H&~sWsuhSBG-ImbnwS->3CG-X@p*L&^ zy-`c(jax!*(h_>pme8BEgq~`N)$7e$!ub|6H>ZEyzuvN2g!8RhLT}v?dYhKe+qQ(> zt|j#LEunX43B6-W=$%?Z@7xl4mzL1GwuIiTCG_qsq4#JBy=P14y;?%=-4c4AmeBjQ zgx;?u^!_cO4`>N}U`yzOT0$S(68eyq(1*5!KCC75+?LRXw}hV868eai&_}j}KB^`3 z(Ji5mX$gI7OX%ZTLLc7}`h=FyC$@w>sU`HuEul|o34Lly=+jz4pWYJsjF!-6wuC;b zCG^=Xq0eaveQrzW^IAfm-xB(Qme3crgubXH^u;Zq=eLAj&=UHRme31ZLSNbvdQnU0 z%UVKT-V*wXme5zWgubdJ^wlk)uW1Q=ZA<9uTH+rq{o`L3IKQEn6Z*!M&^NV&zPTmz zEiIvMZ3%r_OX%BMLf_F6`p%ZnceRARyCw8JEursi34LEn==)njFK(wqtTJjZ!^-v7 zTVkd98!e%~*%JDJme3EjgnpdXtvWo3@1BtR?hROH8hvC)dt%dmIoRABVSup4Srk zh?dYtwuC;aCG^oPp^s?^eQZnU<61%=-xB(Sme41*gg&Vy^vNxuPiYB#YD?(TT0)=R z68emm&}X)UKC314*)5^ZX$gI9OX%}jLZ9Cf`hu3w7q*1Hs3r8pEurVPgkI1R`jVE= z3tK{8+7fzEOX$m5LSNnz`ihp&SGI({swMQ*EupVz34Lu#=<8ZSU*8h?hL+GbwuHW^ zCG^cLp>Ju4f4ubM+WCR*2g0@5gDs&SY6<;tOXx>hLOR9v_Fdgr3(D`iPd$N4A7M zswMQ%EuoKT34Lr!=;K;KAKw!CgqF}JwuC;ZCG^QHp-*WEeQHbS(^^8G-V*wZme6Om zgg&b!^w}+;&uIyLZcFI%T0)=S68eIc&=qQq3>@Ay||qc zG5J39A z`_O$D(~O@7#6MZuuRr>ElJRq8!H@yupayn{_sRo^t&Av(CoPdCvY| z*4bCRaKu-aa^p+)|LOPz?fCipPp9ARN6+W~Wz^&6*Z(x?@pI&V9QF9m^Pg+?&G`QL zceUSc`nPPgHzD*gEuoif3B6oP=;d2Nuh0^D#g@=3wS-=|CG;vSp;v7Qy;@7?)muWZ z(Gq&ime6apgkHNP^qiK^>$HSkw;(0jFn-n%9AJ}sg5Z3(?!OX&SuLVNxgKZoc8YeFB?68hkl(1)~y zKC~tDVJ)HOwuC;sCG@`uLX6C$xkq0ep!eNId0b6Y~6*An{tme3cpgubvP z^hGV9FK!7vza{j7me7~9gkIPZ`qGxri(2AO>dRWnU)>V= znwHSlwuHW}CH}a+z9s%ceM3ul{To|C-_#QN=9c)Q`j(dX_k+HEt*;#LXZ5YUobdLy zwS>OCCG;IFq3>)7eOF89yIVrv(-QjLmeBXLgucHe^x}3(#J?T2m*L;kUvG(jU4Nq` z^fy~VKhP5T!IsbuwS<1SCG;aLp&xAt{a8!r$6G=_(GvQ}me5bNgnqgu^fN7?pKS^K zTubQZTSC9k68goK&@Z)wez_&|D=nd4Z3+EaOX$~ILch@x`puTmZ?(j~s=w6||Dt}o zCGik8qsmrz-xB(Ume4o0gubaI z^vx~t$kIRlb(Zs6dpV(RYYBaOOXxdVLf_dE`mUDHcejMTrzP~gEursg34MP{=*8`n zh!;lfWq7{+dP_W4f1@SzH(NqK&=UH=me3Eignqas^dl{yA8iT!SWD=~TS7n468g!O z&`-65e!3;}GcBQ?Z3+FHovxXFzT1TJ7g|EU*b@4sme4P^gnp$Z^s6nQUuy~ddQ0dx zT0+0s68f!{c((plOFUh_-4gnpmeB9Egnq9j^!qKLKWGX4VN2+bT0(!^68e*t(4V%1 z{;VaQslVM4&cD+V`nxUhRQCZWPX4aYCYy8~y?5wlP$G=(o+^n-( zdfE8Cd4ATJH*<~`MrVIKlha>^Oy9-zTb(V{656XzzsabVX$j}cwuD};CG_$wp;u@L zy<$t~m0Cit+!A_~me8xVgkG&B^y)34*Jug7W=rU`T0*bg5_(Qc=yh5`uiFxOy_V4H zw}jrHCG>_Zp*Ly?y>UzEO4+7f!Rme5lzp*L>{y=6=2ty)5F-4c46meAX_gx;GcW4Q{V@v3rT0-yK5_*@G(7U#T-mN9{?(P4>-d#Xxb{=+`k8LqCLtAE+nJs2! zW@ctg9EaIqRN^=cP8@P@S7wWu!D6&!SenJ>SN+}SRh>HCnX{QQJ61hsp8nsbZdI$h zy6U^H+j54`?`jGC?v~K+X$k$_meB8O3H|<-&>v_C{lS*dA8HBx;g--JX$k$&me3z- z3B6TI=&f5qZ_^Tb+m_JVwS?ZjCG-w0p?7Qvy;Do*om)ch(h_>tme9Mkgx0*p|@8wS+#tCG-g`p-*fHeNs#4lUqWc(h~aAme8lQgg(6` z^cgLo&uj^OR!iu!TSA}H68hYh(C4*;KEEaO1uda3Yzci)OX!PRLSNDn`qGxrm$ihx zye0G%EupV$34K*d=xf_0A|9Hxx8cG1^DXf}{e_m$Uu+5erIyfNZVCOBme5~q3H`N} z(ATwuzP=^&4K1N>Yzci+OX!f=(}4&-_sKM z-j>k!wS>OECG-P!u9o(LT@&s<)DrsPmbkzEdQ03}KhhHV(U#DUwS<1WCG-<5p`UCC z{Zvcnr&~fl(-QjGme9|&gnqsy?yJAi67GMqCG-m|aZmlNmbknAc1!5*w1oa{OWala z&lhW$xO%4R9lQ5~Eunv?CG^KzLjQ0}=ufnS{*ji@KiU%d$67-FcuVM?XbJt3Eunv^ zCG=0Xg#MY9&_CM}+UH_x6Iai4?cU4nbMbTiz6rnopKl5M3oW5P*%JB}TSEU*OXy#2 z3H>WAp?|d{^slvq{`Hp7ztIx^ru=v|9(s8 zPq&2rgO<>L*b@4WT0;MEOXxpo3H_%nq5rHU^q;qc{)?8-f7ufHGcBS2swMPaw}k#| zOX$C83H`S%q5rNW^xwCH{)d*(pKA&Ik1e79sU`G3w}k$emeBv&68hg-LVGRP(uD5= z!o9`sg~jiMm)rNk)^lI|FMTfV()YsnT-4(kL4UlDBee4Yp+C_QdORcO@rhZa#$LFFRpNo2Y zF6!~QsK@7`9-oVPd@kzoxv0nIqQ0g-Cqj?UMLj+j^{0A2LXXcy{psFL=<&Iz$LFFR zpNo2YF6!~QsK@7`9-oVPd@kzoxv0nIqW(<3PC|Pfm|hF?_*~TEb5W1aMLj+j_4r)W z<8x7u&qX~x7xnmD)Z=qekIzLtJ{R@)T-4)pQF|^fF>&=w*X}L87QA>y_WAgJe=Qil z7Ov@Qfbe@?d|i0?y)JA$_ox5T*M;5Zy)NkSjG#Z>#}RrwBj```c0!M51U;S+^msqIjG#Z=+X+3M5%hRQ z(Bm0Fk7oouo)PqTM$qFKL62tyJ)RNtct+6U89|R{1U;S+^msCM(jWZZy)JmIgZG)<&%IaHg!b-pA2RD@TEhKhTS70_ z5_xOv|D{n3`tA8QHi`9GZj^wu?@ zw`mEzZA<9wT0(E%5_*T0&^xw--l-+@&Ml#LX$ie+OX%HNLho+pYH9D$HR1lAEur^n z3B7kq=zUs3@7of3zn0MZw}d{RCG>$Up$}>aeQ-?Yp)YC)eQ`_ZOIku-+7kM*me7~CgubFB z^p!24uWAW>ZM#IoZnO3_>{@@mCH(!-7g|Dpu_g4ET0(!hCG=NXLVvX-^w(NKU)K`) z`j*f)w1mE~CG<@#p>J*peM?K|TU$ck))M;mme6;!gub&S^j$5X?`{cwPfO@~TSDL0 z68ip@&=0hPey}Cf2Sq%cUxk++S&OPCO))u z=6-x+-n}h`dyme$clvPev3d8d8SY*8N^$j7)7Q5;AvgcTPj0o+GHU5-tzL1~XZAL| z!mQ7l_42bmZ`RAr`utfhTl;b1M()a!K3 zsn_kAQ?J)Gr(VBnPQ5|boO;8qIrTX&8au-np1DqHK*RZYfinSYfk;@ zuHoxvnf_okr~BjA^y7UYyrz4c?!UHc`1n`#@tV{9*SQ()zq0$9)BV@G8ScNL`nhL8Wh`}J#1_ut}Xxc`53Uvs+uRyV`_|Fip=)BU%( z8Seid-PfG%zunDn|9|hk=5+rZZif5+sr#DK{VluZ*q7Gc|9|WK-`O>%`|s+SQ@^`w zPW_&)IrV$H=G5=&np3~OYfk-vt~vDwyXMqhA361hyXMp%>6%l2v}^eN{IC6fG^hI? z>zY$<)iu2T|I+(wPWQL&np1DnHK*RTYfimg*PMF$t~vD%U32OkyXMq8b-nVN`yI1vx)CYCVsSoa&QyHbk}hWr0h z_cf>cM|TbPzuWzzyXJKNn65eXv0ZcOeIUB)TejNsn6(|Q=i#2r#`D|PJMRQocf%uIrX_+bL#WD=G5nR&8aWwnp0oc zHK)F)YfgP}*PQy2t~vFkU32Qox`sc$@4V{8|Ex6KzkKQKx%=PlzUFlQisAmZy01Ci zzjC<$LiaVN`&V_%sjuyN$*jN8`)j)YeEs>Zsh_LA&^4$2V%MDdOI>s7FL%wUztT0Q z{%Y5p`fFWt>g&4Z)Yo^-sc&#S^o?C}x_?vGociXjIrS}FbLv~W=G3=!&8ctinp5A= zHK)F_YfgPv*PQz9t~vERU32PtyXMsQb&{aDwW`thzg^%Grl>Ld#fyP?*#7UlrU^z(b)E#}=j z%^PRDX5PK&_v7|!=iQsWKHYoWynEw6@AdQUdH=Kg5O0`wZ+g3XZ=82;{O7%C-o5eT z-aPN#_-lE~ynFsR)7Q{j7q~b5yl-3J-t_14_66>p=O>PLEO2l7oNl?mz3FTAoeSKX ze*f=U;NJA>eD?zPoQd$B1@29soA)koZ~8gkx4^yW=Y0PH_onaV4=iwR`aOJbfqT=} z^M@9=H+{}NyuiKbpZAdk?wz#sxBdekop*1G+1t&>X7}blA3phtB_HwnK0P0%o}Ld= zFPZ&ZQ#%7fA3i)E^j5Qf*5{VK%iNFC=j45}Z|oCi_dYP|>GSfQ*)Qb0+5PvG) zhAp8tY6-n@OXy8nLT}m4eq&4MH?@R*b4%#Aw1j?ZOX#v|D{n3`tA8QG{RZHltTS9Ns5_;Q~(A%|y z-o7RD4lSW~Yze(nOX!_jLhsTNde@fFyS0Sgy(RP>Eur^p3B6ZK=)GG)@6!@`-ciTS6bx68hMd z(8sleKE5UN2`%S%)l2{T_;+{zq;3%2|Kyg?r?kWm*Qd6G`=_;pKD{OM87-mDY>9W( zXSKvT>$6+J`=8Si`rMY#=e5L^_4zIFj-mbEMt^^OL2oB~{Dm!{FKP*WaZBh+T0&pi z68f^1(3iJ_zM>`cl`WyKY6*R9yF|p>X6ODCG-s~p>J#neN#*5n_EKP(h~aCme9AggucBc^c^js?`#QuS4-%- zTSDK{68hej(D${3zP}~(11+H+Yzh5POX!DN;?4EfTjGuNBQ2pHZ3+EYOX$a2LO;mca3^KCl&pViMxEdIUu;@_Jup3N8k z-u#87F9(Z%Z{Fu&`Pt{fYw-$Q6MDs#&?~isUb!XoDlMT`Z3(?vOX$^GLa)&hdd-&5 zYqf-4yCw8GEuq(K3B6uR==EDdZ_pBY!5Qi)Zu2v-ykP4|{e= z`2N_dCG_4cq4#MCy>CnC{aQlr-xAv2nG<^P@6G@7|K8l^Ydo9l|6cyyeEN4-&hO8i z*tY*Ha6wD>{an}*`l6Q5-|O$q-M_T=BlKl0p)YRf z@68|S>wxg@fIQj~`mvVKkGF(=q9wF{pM}s*wS<1UCG;~bp`UFD{aj1v=Ud{x)xXoy zf4%kZoH@ICdhQVW(*Hfu$+LUUEO2jn4cV#J5dWUc_;XGl_w2ln`_Sy?GrP_{?z6ME z`S|JQ`1x18^xsw2{l7WAyY}I8KgFz~Ug@m#N+=?Oia>-Bi9*WE+yyyY&bCnyv`ksYYF}S zme3z)3H`yA(7sov?^pfdn$RC<3GGZz=#RC8-l`?^)-9p8X$ie;OX%%dLT}#^dWV+K zJGO-0sU`HzEunX53B7Ae=-paE@7@x6kCxDTwuE-3C$uv?p`Gao?MzST{aQlr-xB(O zme2>bgg&Sx^uaBm4`~T~XiMnBT0$S*68eai&_}j}KB^`3(Ji5mX$gI7OX%ZTLLc7} z`h=FyC$@w>sU`HuEul|o34Lly=+jz4pWYJsjF!-6wuC;bCG^=Xq0eaveQrzW^IAfm z-xB(Qme3crgubXH^u;ZqFKG#VX-nwKT0&pm68egk&{wvEzN#hkwe1oS`!BP!y$$=- zpKpnM>o2r~{$fk$FSUgJa!cs1w1oa@OX#n)gubpN^z|*FZ)gd9V@v3pT0-C468e^w z(6_dPzO5zn?Jc42XbF91OX#~=Lf_pI`kt21_qK$-uO;;TEukN13H@M8=!aTDKim@g z)L(Cjz3N9=LO+#HemYi|G?A~W*{fd_O$im7pe6K%EulAR3B7Sk=uKKeZ`u-ivzE}Cw}f8O68hCG zp|@CiIrrm%{XxH`E5iNPwuF9NOX%0PgnmOy=r^{6ep5^6H@Ad-OH1gtwuF9LOX#<^ zgnma$=q+19zq2LuyIMlOyCw8{T0+0KCG`7RLchNy^aol(f3PL=hgw2^xFz&QT0(!c zCG^KyLT}X)dh3?Z+q8t3B5y0=p9=^@6-}{=a$gBw1nQZCG>7Bp?7Zy zy+=#vJzGNW)e?H|meBjOgx$Fzh#wk7m&EuoKZ34KCK=o4E)pVSii@Gh0HR)e`#bmeA+4gg&<=^m#3z&uACvI`t!Y=IHLYSOXx4Qg#J=X=r6Z~{z^;e zueOB#T1)8bT0&pn68eUg&^NY(zNsbj%`KsCX$gI6OX%BLLf_sJ`i_>+ceaGSt0nZ^ zEurse34L!%==)ki-`^7YftJt@wuFADCG^8Bad`dpmN>M2q$TvDEukN43H^9W=qFl2 zKiLxcsg}@Bw}gJCCG@i`p`U9B{d`LtR)3=<-2Y}v=oebzkosFKad7?ZmeAj63H{xc zIH;a~u0G3~;^^7E&o0fmGxaNG*C%%I%B~6hs+Q0`eD0^1^|CGD{&Fp$mv0HZLQCis zTSBkY5_;v9(5tkBUbQ9kYAvBxZwb9dOXxLQLa)^ldhM3b>$HSkw{y`&}HU%$E~+}~p9<=l^l*RSb{aR0R}pcL-_R2JjV+D-_a6!%a+jZYzh6YmeB8R3H_dy z(C=*t{l1pa?{5kHftJu8Yzh6Lme3z=3H_0l&>w9H{jrwNTeXDVx+U~BEupt<3B6rQ z=LpU@KeL_1eY`=qW3_fKvKeM(E{Q(Hox))M;kme6Okgg&z+^jR%&%rY-M zSG#{sZzuG*EuqhA34MM`=nGmxU)U1*qL$DXw}if=CG@2&p)YF*eR)ghD_TNd*%JDy zmeAL>OGKP7Yj4Bx_2*mSxcUn%p}*J?`b#aLzuXe~D=neF+7kL}EupV#34MJ_=o?x> z-`EoRrk2n*w}if>CG@Q=p>JymeS1sjJ6b~D*%JD$me6;%gubUG^t~;i?`sKte@o~G zT0%eA68fQ*&=0r7vGvzm;^_L3me7y3gnq0g^y4j|pJ)mFWJ~C$T0%eF68f2z(9gDn zey%0-^DS{q{f(Az|C=qLUucP=>Tk8gk@dG*LVu?v^mkk0h}yaOvnJ%+n_AynuWSkZ zs+Q2pw1i%^CG>JFp_gw7y+TXq6UbiLmdM%;XZwb9YOXv+-LT}U(dgGSRo3w=9v?cUrEulAW3B9Bx7SGO$ zXXkDDIv{*~Y}*ogyOz+~w}jrICG?Ihp?7Kty>mwS>OCCG;IFq3>)7eOF89yIVrv(-QjL zmeBXLgucHe^aCxSA8ZNzP)q2CTVnC-ym)qAW$Bl~+-ujWT@zl*R%;2pdQ0dvT0*bc z5_+wc&}+AZUZ*AWx-FsCYYDx6OXv++LT}g-dZU)m8@Ghsq$TvGEulAS3B7qs=p`+& zcy?YqJ8#q10paUo@$9^Kc3wO?Fa3TvuCwrYT@t=O&Tk2QK}+ZhTS8yd68hqn(3iA? zzO*IuWi6pEZwY-xOXw?GLSNMq`r3Ahh{eBmUOYQL($@iTbYCBjwuFAHCG_Jhp`U08 z{bWn%r&>Zk-4gnlme9|(gnq6i^z$w8-|Fw3`(?~=%52T?y;k@;>FM4za>hEZa`EE5 z>Gi~0Pnh>{AM(Z-C(gU)&8N-ooiw}m*;)JNPXFvxyLe?yXupKHU&gF`_}ou1>t$QQ z`!ClLdij>nE3|}Ou_g3MEumL#3GEC&o#XX*j@RQkUXSN^J)YzBc#hZOIbN^bKabGk zIbM(Fcs-uu^>~ig<2hcB=XmW5Pw4R+ug7z|9?$W5Jjd(t9Iu!3^AfAquWkwVw^({P z_u~ia*K|d=|Js(&uWJeY`j*gdXbJtsme6l%3H|1l&~Ir8{nnPyZ)*wt_Lk7^XbHV# zOXzpDgnm~`=y$h-eoss2_qK$7UrXrsw}k#cOXv@_g!X+reJ|?|*M$B^OX!cbg#K7d z=&f2pZ`~4lo0ibqwuIiUCG_?!p?7Ery<d zXG`e4T0-yL5_+GO(EGN8-mfL}{w<*oXbF8_OX!1ILLb}``jD2;hqi=1tR?i}EuoKS z34LTs=%ZRfAKeoAn3mATwuC;eCG_zvvGJ>3`gh>&pV-?8eNs#4lUqWc(h~aAme8lQ zgg(6`^cgLo&uj^OR!iu!TSA}H68hYh(C4*;KEEaO1uda3Yzci)OX!PRLSNDn`qGxr zm$ihxye0G%EupV$34K*d=xf_0B37BTw_)Y_^DVJb{e_m$Uu+5erIyfNZVCOBme5~q z3H`N}(ATwuzP=^&4K1N>Yzci+OX!f=(}4& z-_sKM-j>k!wS>OECG-O=p&x7s{ZLEjhg)LB`s*#ReEmpE=to;ZKh_fZ@s`j}w1j@L zCG=A*p`UIE{Y*>fXInx)*An{qmROyDhOy z?L7Wj6LRiNt?#W@wuF9FOXy`Zgg&q(^g%754{iy4 zNK5ELTS6b!68i9#&_}d{KC&hBQ7xg5ZV7!%OXy=;LLb)>`uLX6C$xmK^ocFu{z)yN zPi_f)N=xWdTSA}K68iL(&}X!SKC>nCSuLT@ZV7!(OXzc3LZ8Fa=S)_k-j^kXfd zA8!f$L`&!=TS7n868hhfmKh+ZYr&~h*OiSpWZ3*qqd216_&vfnH%k6XVbN#*vzyF_a3H=K#p+DIY`WIV5 z|58inUv3HgD=ndawI%egwS@llme9Y^68bk=LjP7v=-+M${W~q8uW1SWyDg!AuO;-S zT0;MROXyFxg#Lq;(0|wx`j6~fE$u(Jw6xp_*~TEb5VOPE-`WSOxNx$z81WA zM)vvmet#_(zZS0PYk=^3UwmD7`MoY|J@=>o($|H*nD@G%$1{TdcppdT@r~YQ9k4t~>d-uBFwP@mm*~LH1-FxxBU;f8g zpE?6Jby`BN+Y)-cmeA|Bgx;Vf^oA{= zH);vJaZBh;T0(Ey5_+?i(3`h}UeXf!)h(g7Sb914<5~Sdzosj~{nxgHeqBrG*SCaz zLrds4wuF9DOXxSZgnmm)=(o0nep^fEx3`3TM@#4}TSC9HCG@*mLchBu^m|%Dzqcjy z`&vT3za{hsT0(!YCG>|{LVvg=^ha7kf3zj^$67)=PfTYDy>(6KZCXNa+Y)-ameAX` zgx;Yg^o}i|cWMc}b4%!5T0-yI5_-3m(7U&U-lHY-o-LvGY6-n}OXz)CLhsuWdcT&? z`?rKXpe6KyEujx;34L%&=tEjUAKDW7u$It=w}d{TCG?Rkp^s__eRNCcV_HHV+Y$B|cl9)DqtR@Gh5=X>$6(ouj;c~ z!uy}o68hYh(C4+pXX^7?;xC8x{|D;l>I-^1;o~oC34KvZ=!;uIU(yo#(w5MdwS>OB zCG-_7p|5NSeN{{7YuhCv{(RQnhCi!6-xB_Q>I*HQzt|G`OD&uH{{X$FpVg0R^_=Eb}Eup{D68gI>@#%W; zZ0>Wg^tEJ_mhjrLYD?(VT0*bh5_*l6&}+7YUaKYa+AX2iX$ieE+yyi)Zu2v-x&?B@n(owr>f& zLrdr#TSD*D5_;#B(7UvR-nAw4ZY`mAZwb9eOXxjYLhscQdheFd`?Q4KwBU?fr)e`#Xme9wvgg&+<^l>er zk8cTmLQDLYp3To$=H-1hKfB-m#eLogpU+EL;x+Z9E#dxUEuk-O34KLN=qp=7U)2)& z+IESE#k2Y1+5C~d4hU!2M_WQa))M;hme7l5^QE7Q#lJUSJezOYmjdCncC(hyo415s z(h~aBEuptqdO7#w;@N!hY`$Gz352hY?OQ_c&=Pv_Z2s~+o1eGL|F^$4KevxRwXY+> z*VAb&p-*oKeMU>@Gh1Tu@68v_<~Q|qK=^y(n_EKP(h~aCme9AggucBc^c^js?`#Qu zS4-%-TSDK{68hej(D${3zP}~(11+H+Yzh5POX!DNV)5_I7ysUT`K4b9bFZB%bWM0I zU9lzfN-d#RZVA0gOXyWwLa){mdi9pjYqW%3vnBLeEuq(L3B68B=yh8{uh$ZK{g%)h zw1nQUCGx#k2Y1zti&3rC$znUk@Man(+0pRZHltTS9Ns5_;Q~(A%|y zUObyG{=NC}OTQfEz8+5Kn)omMdvjk;r}S}z{*OGHpWj*ktbSf%@!x4#{Co4ov-#rL z{Dq}22aEqsOP`0uv-#352G5?wv-#q`Z~lSVYlE+c4|YxH54D8;a7*Zqw1jppUp$-7 zeLChoAA5F5`26hE5_<2J(EGH6-nS+6el4N*ZwY-sOXvezLLbx;`rww(hqQz~v?cUm zEujx@34KIM=p$P~AJr22=$6pOw1hskCG>GEp`FYBYn{!Vo&O`x=8J!CzWDd%H}&;O z_9q{gYum+h{o?<>RpY%&XZOaR=j6GM?tk|^`*|*!_2u8==ec6uy-S9F z{*?>-Gp1kL)w6%bS+jq}Rr7wH<7aO-KREB+^xw(4_e1mUdD{{{#K-5|yLjoF=YIU* zdH1H@tDB#gcW?T868C;&-o5GXRNVW~dH2ru#u-00@80xzaQnyS-J3oq?)}8Pd()q_ zdp|kv-i6*c`1}5u*}c!s-v1RXaaR|OXc-nb?7CM}^iZ3(?uOX$s8Lch8t^cG7m=YCwL zKX^U8&HdNZgnn&H=-0J`etk>mH?)L)V@v2awS<0iOX#<>gnnyF=(n|metS#kceI4w zvL*C8TSC99CG@*nLcgaa^m|)Ezpo|q`&&Z$&xF%|F4P~a3H_m#&>wCI{gIZ?A8iT! zv6j$VwS?ZfCG<8ep|@=by0))M;g zme5DEgg&w*^ieILk8TNlOiSovTS6b#68iX-&?mHnKCvbANiCsIZV7!#OXyQuLZ8+W z`t+92XS9SqvnBLdEuqhD34KmW=yO{_pVtz9_sW<4-|uw)g5FN(3tK{8)DrsQme7~9 zgub*T^kprfFK-EbMN8-_TS8ye68hS9iHJLA?QOWD{(MW^UVot_^cPz~f2k$(ms>)A zr6u%NTS9-WCG>SIp|5WVeM3v=8(Tu()DrsUme9Afgub;U^ldGnZ?|)`wD0JeaR1Jh z(08?jzPly#JuRW{Z3%r}OX&MsLO;+F`oWga54D7TxFv3@zupqJ)Q_};ezYa@V=bW{ zZwdWGOXw$CLO<0K`stR?&$NVowk7m)Euo)piCgP$w1oTLYzh5BOWa(4t0iu#zugl0 zJ1wEV+Y&d{&i8ZE)p>HpXJ6^!#d{}BH`-?p8_s;JpoPLkLJn!BK(?6{Jm3jA0>1|79_kVTXJ#U`w{o3r_ z8ngGkx+U&?)l2{XEYUyI+X?;gme4=k68aM@p?{<$^pCcL{;`(OKi(4hCt5=PWJ~Cu zY6<<*Eunv=CG^j>g!c8bwTY`|x_0jeTS9-oe=dHm-zVYs`|~ZKf1xGxCtE`QVoT^> zY6<z2@;Z3+E1EusImCG_95g#P=M(Erd9`g1Lz|FI?XKedGZ=a$g_(h~Y#TSEU^OK7hJ zTbl5FK)AQ~y|DPb@N)ZJ*m~}(|E15x2l`$ZpNo2YF6xi>afEh0AoM3%LXT$zJ)RNt zct+6U89|R{1U;S+^mshZa#$LFFRpNo2YE^5!kB_^((>Ds-;*Mb+%$UYz6 z@2>^p*TOY@4G@0si?0hWzt@GW=l=9x`nvG&yw?Rio)PrN`#3_6X9WF;-cIQ8jG)Ie zf*#KZdORcO@riq~ko)Pq?dpn`W zGlCw^2zopt=<$r8$1{Q+&j@-vBk1vrpvN0q-grjPi?0PQ{ymz{ z$M@T7!Q%J+i{}7;zAuOG{a5&Ak89?BT>69GyVnJ;b?`pZ`?>eZn$X^T?n7q1OiQ@G zY)k0nT0$@15_*M}&?~ltUa2MY$}OQ+X$ie*OX$^FLa*KudX1LQYqo@5t0nZ>Euq(G z3B7Jh==EAcuip}SgO<=6wuIiOCG^HEp*Lv>y=hD6&00cl-V%CAOXyd(gx+H5<=l_I z-XHX9x+2_vZA<9awS<0sOXxSWgnna7=r^^5esfFcx3q+QYfI?2wS<0qOXzpBgx<0x z^gCNZzpEwmyIVrPrzP}zTSC9DCG`7SLVuto^aoo)f2bw&hg(8_q$Tu6TS9-VCA8=N zbOz8{*M#1tCG@r}p|@)Zy?smQ9a=)~*b;iDme4!5gx;kk^sX(TcWVi~drRm&T0-yH z5_+$e(0jLp-lrw>zAd5mYYDx7OXveyLLb-?`k z=qp-6U)d7+s+Q2#wo61jG;43egZ1ZI!rvc#p(XSeTS9-SCG?kDLVu+t^jBL#f2}3- zbuFQqQq3>@A{eYdTrTt*ng!>P*gnqas?ytYz68<~XM_NKZ+7kM)me7y4gnptW^ph>2 zpK1yHbW7-GT0%eD68gE8(9gHTef2k5!u@ZygnpqV?y0}k5_i|%ZVCOJmeAjAiMwiN z=jF|eCA0nQWnMhTded*#E6wgrz3RMPvOrI7`>j5H`o8pTzft>Vzrw_lxgVGQ@H?}6 z(b+|OD~!C{-+Fm&AeaN@6P*m9W}h~@6GNlH@xkb+4b*t zZ~S#VKJVA1$6uFTGW=pSnd{o^g6f1)MyPqu{qsg}?`-4gm|T0;M9OK4v!TbsCgrfc_JZl8;v>-SCg z{r`MR=wE0F{mGWlzt|G`ms&#qa!cr6X$k$SEunv{CG@Yig#L||(7)Ld`nOs_|8`60 z-)RYbO-tzCZ3+E*EulZv68iUBLVvm?^dGc@{==5gf7BBCk6S|jNlWNIZ3+EnEusIs zCG=mkg#OEx(4T1u{Z}oa|GFjgXIn!5O-ty%Z3+E%EusIuCGelIM3FTC8o7q*`J>VN5T@wvVi#^<6Q&j|YC zeH@{k4+#Bh{W%eOd@kzo zxu`$Y`w@D4F6vMBc0!NOMLj+j_4r)W<8x7u&qX~x7xnmD)Z=qekIzLtJ{R?8`gIc8 z>%jC{pvUK;9-oVPd@kzoxv0nIq8^`%dVDVG@wuqS=b|2;i+X%6>hZa#$LFFRpNraa zafykmXS#N8@wMQ^GqTUe_xo$X__c6NUju~S`{L`u%kOnz>$yMum%c81bKdKM9?uB+ z<9!^V$1{TdL~kebct+6U89|R{1U;S+^mswE&KUv`3 zakG1Wy1>12X7~PVfqN&+?)~}qxc3(e+&gdfaeukMz2j&1KJz{9{nY~Zrr*O~FL3Yd z+0Xge1@0X)yZ1K>+?#&yf4jiF>Er%xfqT6+{>%hIQ&AT`K-rT!x-o5Gb?%wtD?j7%qGj5o7Z~C)!`^I_q&X_)^ zy=mUP>Cel%+&u5z^u6QWE%WYq+w|vf>jL+V_YQH}0{5mrquUp_H~l{DSm56DT6O0F z_ohF?JLbMQUizAP{L){_58SoD$4&R{o_BBhz4_34=H2tQ>ErJG9{28>cW-*#_wo17 zyEp!v4=iwR`q=ruzwyuZ!393U;L{7-o4)R!S>WFE_4(|) zd++i|n=${-=zY_TcK$uLrl0eEuq(G3B7Jh==EAcuip}SgO<=6wuIiOCG^HEp*Lv> zy=hD6&00cl-V%CAOZ-3V-35@P=V7P&$d)W-=8?~EujAC0XS!8A}-Ku-*KELiht5;opeQV3PH@nt(D$o4;nQyvh=FI8yz0)2$ zdXtvWZ*2+vwwBPFwS<0qOXzpBgxW$RZHltTS9Ns5_;Q~(A%|y-o7QY z&%Md#pWdM+^o}i|cWMc}b4zHi0h8AOy=zVA-C9EL-V%C`me6~)gx;$q^xiF@_h||3 zYZ#NSW9a>ALhs)a`hb?u2eyPhs3r8lEujx-34Lfw=)+nXc}wUkT0&ph z68fr^&{wyFzNRJgwJo8qYYBaQOXwR~Lf_aD`lgo9H@Aeor6u&OEun8~34MD@=sQ|M z-`NuSu9nbuw}if@CG-RBA`$<6u6-Z=Tm6xi_^*?I2zjdRz3 zqw7obHnDBj{gAmIHq&pe3D^BnXMVYJy;MuM{+5=|OSgnxrX}>UEusCA=YH8Uy?jl$ zzCugr6UbiLmdM%;X zZwb9YOXv+-LT}U(dgGSRi&{c&(h~ZuEur7m5_+?i&~I-E{f?H6qgz5B(-QjFme9wwgg(9{^a(AYPizT&QcLKQTSA}G68hAZ(5JP8 zKD{OM87-mDYzci{wsY z5`O%pEuk-K34M7>=qp-6U)d7+s+Q1Kx5N(hH7)VZ`r4N8wFJ{fU;)pKJ;Jsg}^6ZVCOFme8MV3H`a2(4TJ! z{a{Pzhgw2E+!Fecme7y3gnq0g^y4j|zt9r;i!Gs_XbJsfOX#OsLO*sEup{D67Q(L+Y)cDztdQnU0O~wI%f1T0(Eu68h~eq2JLGdh?dhTeO7UvL*CZEupt=3B65A z=xyyx&GvS^O}M^&OXzpDgx;Yg^o}i|cWMc}b4%!5T0-yI5_-3m(7U&U-lHY-o-LvG zY6-n}OXz)CLhsuWdcT&?`?rKXpe6KyEujx;34L%&=tEjUAKDW7u$It=w}d{TCG?Rk zp^s__eRNCcV_HHV+YEuq(G3B7Jh==EAcuip}SgO<=6 zwuIiOCG^HEp%=A;-lQe;TU$cEttIqkEup7p@#$H7dKQ1}^WoIqBz%6H))M;kme6Ok zgg&z+^jR&T&u$5QPD|)>TSA}L68ij>&=<6XzOW_qKj&He?mpgSeO$uFeC*Tp~6<3Kn|ezqm_=UPI4z9sa7EukN33H@+O=to*YKiU#{dKRC4 zUHk{LPlx~Ub@4t;XSiW*%@Q+*)Yty|6gLjn-Zbx;OOxyW*OKbB_fwbpCf>Zj`%SLR z9+DS+za{6r-!-!zZsyyoXAax1z2BnY+U!4YS@_x#OZkEK-ZMEn->VYG_Ivorme4=d z68c>&q2JvS`aLb7f4U{~&$NX8*_O~h*An{YTSEUrOXy#03H?hgp?|q0^slsp_E_v^ z;@&yE?b;jdvG~<~--O@)ueF5!S1qA`y(RR&ZVCNwT0;LuOX%Nh3H@6wp?|w2^uKKh z{W~q8f43#{@3n;f{g%*w&=UH-me7CL68eu?Lcg~q^!r*u-`^7Yk6S|jNlWNIZ3+En zEusIsCG=mkg#OEx(0|nu`mb9;|4mEi_qT-p+m_HDXbJtnme7CK68i63LjOZc=znYp z{h^l7A8rZ#Pc5PUT}$YH-xB&iw1obTEusHYOK7hJ+nextK)5#jT$p|?ywN@v_M3V3 zzy7{>LZ1uceNm5R1U;S+^t<~wgnmy;=<$r8$1{Q+&j@-vBk1vrpvN-zB`=TE2i+a2->hZp)$NQol?~8i8FY58WsK@)F9`B2K zyf5nUzNp9hq8{&y+I?}6iF@btwrkVZg4fQ-9*-aQ*Mjl6a9_^>!tZ_hy70z(UD$8t zPyh9=3%~Z}*S{|4@r7(Bm0Fk7oouo)PqTM$qFKL62tyJ)RNtct+6U89|R{1U;S+ z^msu4U2AGgF54y=Kz1cZ-meNC492S zeKX(Ae&hGfzd(mU{iQMZ>i>&s|$(?)&a-i8Gdb{p*PyUr+SA`f-GQcT4E^ zw1ghd7J58e=<#f!$FqeV&lY+-Tj=p@p~tg@9?uqfJX>gw;eICWozvT{jjyfx$NjMw z&sl!I<2g%@=PW(Gw(9YmrN?uY9?w~NJZI_goTbNemLAVpdOT<8@tmc{bCw>@S$aHY z>HGR~C-nGQtH*Pees4dH(C=#reSb^n@tmc{bCw>@S$aHY>G7PU$8(k*&slmrXX){r zrN?uY9?w~NJZI_gwN{VkEIppH^mxwF<2g%@=PW&*v-Eh*(&ITxkLN5sp0o6L&eG#K zOONL)?X_Tg6VuOy>F0va2Y7D05kD99oB7?p{=Rs2kHvUj)Z=|ok7oq^?miBo-_sI$ zJR|7wjG)Ief*#KZdORcO@rD1f zcwf}xeNm71MLpgZ^>|;@<9$(&_eDM47xj2w)Z=|o-`AfLp~w589`B2Kyf5na^?roD zza{i|U)1A$QIGdUJ>D1fcwf}xeNm71MLpgZ^>|;@<9$(k9hkfp=<&5ckM~7A-WTD1fcwf}xeNm71MLpgZ^>|;@<9$(&_eDM47q$E1A`|z{>224huLZNO z2b0I+$NjZnd@kJAZ<_F2n7%H&_O%I*#~a~wVZWI_{nx)P{Q8?;|GJ>ZGlG6s??>o& zw}gIAOX%^8pvN^msmu(5XTubQXTSBkU z5_-j!&?~isUb!XoDlMT`Z3(?vOX$^GLa)&hdd-&5Yqf-4yCw8GEuq(K3B6uR==EDd zZ_pBY!vnBKnEunX83B6NG=$%_a@6r-_*Ot(` zwS?ZiCG;LGq4#VFy;n=(-PX*U~)dt`_+Wrza{hmEujx=34KsY=!08AAJP)~ z(3a4LwS+#rCG-(3p^t0{eN;>6qgz5B(-QjFme9wwgg(9{^a(AYPizT&QcLKQTSA}G z68hAZ(5JP8KD{OM87-mDYzci=9eNRj12iip<7SFZs z!$tK+TH?a`qb;F7))M;TEulZr68e)Zp+D6U`qM3;KhqNWvn`=N*An{kEukN53H?w@ z=!aWEKhhHV(U#DUwS<1WCG;0sLVvL(^b;+ipKJ;JR7>cmTS7n668hPe(9gAme!eC2 z3oUU${iT*TuYR#5^h+(FUv3HgN=xWhTS9-iCG=NXLVvX-^w(NKf4wF2H(Ek}vn9^2 zzts}1f4e30cUt1y`nxT0PW`=>(BE$f{ezY`yPlq%JqEL{C9AfC*Ot{Euq(G3B7Jh==EAcuip}SgO<=6wuIiOCG^HEp%=Bp^z1x6JMZ7~fbje{ zpe6KyEujx;34L%&=tEjUAKDW7u$It=w}d{TCG?Rkp^s__eRNCcV_HHV+Y0jX8V%fCS1R?CG=%2p)YRV^mQ$vuWt!`LrdrzTjDKCP0!9B?{OfU)jrV@`jai8Kh+ZY(=DMt(-Qi#EulZx z68iHkp&x7s{ZLEjhg(8F(h~a7me7y2gnqmw^cPw}f3YR>6D^^iYzh5TOX#OtLO;_I z`q`Gy&$Wbpz9sYvEipYiPtVRP%sv%nUdvYOZNh8YN-d#RZVA0gOXyWwLa){mdi9pj zYqW%3vnBLeEuq(L3B68B=yh8{uh$ZK{g%)hw1nQUCG^wa?PtVS;eLmbU`{vBIK0j{kEkfVa68h$r(6_XNzO^OvZ7rd1ZwY-zOXxdW zLf_RA`tFv{_q2q5pj{+l`gP~&+4;qu2gG?jKVE7H{c=m_S6V{9+7kN9Eup{C68ftx zp}*D=`s*#BztIxycd88cixw`S@79zFhBk3HlZ%gkMS%UmxzxzJv`Xa4BTzP)?U^)fY~J+_m_SNHJzkAFj> zm#+!!?=h3V%V>}N~wI%f1T0(Eu68h~e zq2JLGdh?dhTeO7UvL*CZEupt=3B65A=xtjnWOX$5@LhsWOdf%4N`?ZALza{hmEujx=34KsY z=!08AAJP)~(3a4LwS+#rCG-(3p^t0{eN;>6qgz5B(-QjFme9wwgg(9{^a(AYPizT& zQcLKQTSA}G68hAZ(5JP8KD{OM87-mDYzci=9eNRj12iip<)|qSH zhqddEw8UEVM_WRFtR?ivTS9-LCG;m-LVv0y^ru@wf2Jk$XInylt|j#6TS7nB68fQ* z&=0qSexxPzqb;EyYYF{$OXx4Og#Kbn=qFl2KiLxcsg}@Bw}gJCCG@i`p`U9B{d`O4 z7g}P?`b#abdi`Qc=$BeTzuXe~m6p)2wuJt2OX#n(g#Kzv=&!Ye{(4L3Z?uH|W=pJ5 zf2$>2|8`60@3h2f^>wlQ*8|M1&=X&z7|FK@8Ys6pIZ)yqcT?oBo zOK88mnP1{uzojKyU%DmqGA*H(Z3(?xOX%fWLa)#gdc~H|E474Pxh3=}EumL!3B6iN z=+#?7uh9~E&6dz>wS-=~CG2EunX93B6lO=-pdF@6i%^&z8`8wS?ZgCGagg&?>^dT*w4{Zs3SWD=`TS6bv68gxN&_}g|KDs6J zF)g8wZ3%r`OX%ZULZ8qQ`oxydC$)q=xh3=|Eul|s34K~i=+j$5pV1Qf%$Cq+wS+#q zCGme3cqgkIbd8!h$v^Zf_Aeo5a?`0bjw}k#oOX$zGg#KJh=+C!=ey}CiLO<3L`tg>~UuX&a#g@=dw1j@LCG=A*p`UIE{Y*>fXInx)*An{qme4P>#Glt+ zYKcFsUu+5eQcLKUTSC9m68hDa&|hu|{gsx`Uu_BfwU*FdZwdX4meAj9i9f5q)e^3M zyCw8@TH;UY@3zDr*WYUi{r#5EKWK^jYv=tBn>iNG?euIuJ)3XXbAj+$yHQK%jax!5 zY6-nbOXy8!-<bZwY-uOXw3@LZ8$U`s9|-r?iAVwI%dvEul|u34KOO=rdbF zpVboj?3U2yw1hslCG>eMq0es#eL+j;3tK{8)DqggXmUULyZYjq_+Ra8{PYv&4mP7q#8S8NHrQcLKS zTSBkW5_;8^(5tnCUcDvs8ZDvMYze(qOX#&*La);jdfk@L>$QYlza{hrEup7p^Xb>k z_nv(^%se0V>21REW8aq0`?ZALza{hmEujx=34KsY=!08AAJP(fdN!Y)%`cgKI?OyD zF70i?^W(CX(3iJ_zM>`cl`WyKY6*RHOU(be`Ty+M+{e42k3;-_{dMz2OHIFSK0TXH zzi$4u*-y8buXBIBw~6W3&8PpSWqLNB{al!S-F*6W^8OA9cG>n zi+h{+QNM0}%aT9dXY)-ud*9OIMEHEVwI%dzEun9334KRP=sR0N-_;WO?v~K^w1j@3 zT_j@qZ(62bH-E9`0pTqCQcLKUTSC9m68hDa&|hu|{gsx`Uu_BPTs}RU_y13#$6|>d z3&Q8>n_5Eme=|i&>*jvHxAprb^zAL7?`R2qXG`e2T0&3H=F@-E^8D=6Vdmf9 zywKak^xvCL|GoLrvrmPY*NSC&oABDPY)k0nT0$@15_vnBKnEunX83B6NGXy@|j*?i{Vn0Y*=|EA@Qdp7@3zi#gH>7Vm# zeq&!qx~8Wl@&EU~X?bgZSGl>Lm+iK*_zZc?1{rBc?>#ROKn@|6{pij&`oo3F?Pxdz9EbUxAJ)6%w95ZKW z|JI!FZ(hF968f7hG5z0Znf{xWHD;d*GiQ-Cdzx4{Xn}&#Pr{rPyfC7i#-nrXWy4vLciP+`jwW@{;fHo-w6NSe2LlL z?##E7{|xNyz1I2PSA5fO?ZUaY-!XITwf{VA^SK{8ew?f4e(V+tT${XJZn?m<#dGhs z)dJTpnY*_2ylb1zeZSde-ZkHMlW$_%xohL^Y4UZL?PspN_P=qp-Mo*pc=kigeEZIM z*CzkDri(kwyLQ#^eml;)cH?kur+L@LkNFEr`hoY}GkJU$5AV10y!RWgeXQs2_BjDUEg)ywaMe>+HUi%jgS4d^M1bjyu%!yn0IaR9B_H}c^`9p%=VafZSmv- zwtLRIcK&c}uLZ76uJ66TwaN2fp9QW>KJ30fdyU%yd@zT7{g~`wD@{vE}{a%^(et!JqHRPyy@8{a& zbL{9Ja_yJ}t}XV?aqJJdcH9EjCO_Zt3tXE#HYY4_ZSsC6E^uw~+H=wZ*Cu~XCogbq z@_BU10@tpb{q6n4-_Cok`tHf=?x_pB-{ke;(|V?Hu?4YG1tz!HvW9q&AT>v zOuXOq^R7)kXI;Bt-ZkGhdCuQB@7m;Z$d9>c-nGg5xOVfrYvcRfGVhunKlyXGb%AS( z{eZY_foqfd!tD!On|yBGvB0&-ujkGMu1y}(yMD;EyBD}N`E$BwfoqeGdG8Oo_LB=- zn|#clTHxB`W4>#FYm@hT_X5`@@AsYsu1((Wrx&<3dB2}o;M(NR@9~+N@@r>Ze{Pf4 zx}RO({U(1#3%%cB?;Jn3!23-;{}%e3p1j}BFYtbo=gTiFaBcGZ{ly=0?Uxp~ws`iJ z`4hjqz_rP1-LK5Mw)@=snEAi=oqou*oqxzRmnM(Pug?8^tIYlQdt2fMOa8=9{A5e$ zpK1yHu9nd6ZVCOKme4=l68dLaLjP<_=$~r|{qrrMf1xGxFSdmKrIygY+!Fd%T0%Rw z?q}lOIlb-L8||_9)qdZE-~X?*g#K48p?|$4^uKNi{cl=A|3*vb-)sr}TP>k~yCw9$ zZ3+E5Eunw6CG_vLg#P`O(0|Yp`o5OXf7lZGk6J>%w&>v_C{lS*df7cTF?^{CuLrds? zYzh6Lme3z=3H?tkq5oY==zrf5`aiUU{*Nu8|5Hn7uLawi@OePEHvL?felEPxJ{R_z zdG^2lzW87JTo~_*dORcO@r7xj2w)Z=|okM~7A-WTD1fcwf}xeNm71 zMLpgZ^>|;@_x0yQ=<&X&$NQol?~D3D1fcwf}-@7GCauLG0U0zKXr^>|;@<9$(&_eDM47xj2w)Z=|okM~7A-WTD1fcwf}+i;GO$JEyl@o4yvjc1HGi{J6gsjL(JpdJYhN@6*?XH{R>Q zelvgiuYX>^FYzUKhO9!H=2zIM?1(6Z*|9p}jkymud<9mX^>N}U`yzOT0$S(68eyq z(1*5!KCC75;Vq$$XbF8}OX#CoLLc1{`k0o`$F_t%t|j#GEul|n34LNq=#yGPpWG7q zl$Ow^wuC;dCG_bnq0eXueP&DOvsyx*-4gnomeA+6gg&n&^!Y8JFK7vUVN2+XT0$>w z34L)(=u28cU)mD-vX;=7w}if;CG?dop|5HQeRWIdYg$5I+YX+V|nV)gNgI zUqAY2OX!cag#LI-=ufnS{$xw&Pql>pbW7;Zw1oa_OX$zFg#LU>=m%RuKhzTX;g-;k zw1j@NCG=x0p&xGv{e_m$Uu+5eL`&!=TS7n868h{jHX8 z{o5^}zta-mufN+8->bi8XKJ>;-`j-iKWK^X*3QmLn7DCvX1={%qnJay* zYTh;PH+kRRnY;GExgYECnLI`p_r{xQLch5sw1;)(p`B|VZstSI^;=rPk6*ea^fE1> zJ)|=a>s&9_60UpbW*)w|UZEviU$G_hN-d#RZVA0gOXyWwLa){mdi9pjYqW%3vnBLe zEuq(L3B68B=yh8{uh$ZK{g%)hw1nQUCGCy;RDYx;F04P=68d8;p+DXd`V%doKiLxcQ!SxC-4gmUEulZ#68dv3 zp+Dae`oWga54D7TxFz%>EukN63H?}0=*L?^f1xGx7h6I<(GvQ}me5bNgnqgu^fN7? zpKS^KTubQZTSC9k5*O58YKime7h6KV)DrsTme8-XgnqRp^p{&gf2Ae#S6f1VttIr= zTS9-MCG6XyTw1i%^CG>JFp_gw7y+TXq6UbiLmdM%;XZwb9YOXv+-LT}U(dgGSRi&|oO zcAlP{_wRW?czztv68gZF&LpU@Ke#Fo$}wS+#oCG;sRp-*iIeOgQC(_2EH(GvR1me6Olgg(0^ z^f@h|&us~PUQ6inTS8yZ68gfH&=<9YUfdGS(igXc>zA~IzO*IuWi6pEZwY-xOXw?G zLSNMq`s$X@*R+Jbwk7m+EupV(34KFL=o?!?-_#QN=9bX6w1mF3CG>4Ap>J;qeMd{^ zJ6l5E)e`#dmeBXKgnpo1Bw~7Yo}Qgw?0G<(*Yo40me4P^gnp$Z^s6nQzuXe~D=neF z+7kL}Eup{O68al0p}*M@|59h?ei<`NKJUfvF6GJ{hpqItdP z(7)f0_jx@{b}#_rkc=iZVBz(XFkMS`>6=m z-_jC#>6XyTw1oCJOdbpEahN<7dik2rE3}06I87cay;4o+m0Lou(h_>rme8xUgkHTR z^cpRp*K7&BR!iu$TSBkX5_;X1(Cf8?UcV*u1}&jCYze(lOX!VTLN96wy-7>xx3+|S zTTAH8T0+0QCGJvleRE6bTUtWj+7kM6D^^iYzh5TOX#OtLO;_I`q`Gy z&$Wbpz9sYvE%CnkOD*w7^@}Z`Uup^ca!cq}T0+0t68g(6p}*1+`l~IWzt$4^>n)+b z(GvQbE%Dy^TP@-Gw_8GgrzQTd{%%X$SAVZ1^!Hmr|DYxQpmtVX!o+$&&W9 zYm?7m-}`6t-fvUiIme&RyS8X@q5X?_*A@@g{&McxDs$g=Z%eG+wV!MW{ZlQW-_;WO z-7TTt(-QipTSEU#OX#0%3H@^|p?|(5^e?o8{>7Hiztj@?ms>*rN=s;u{eC9yozvT{ zz0n?vU+woz`2GJ{OXz>q68hI$LjUWQ(Ep|-^l!9;{>_%qzts}@w_8I0+m_J3(-Qi3 zTSEU{OX%Nk3H=8xq3>%6{f8}~|EMMOds{-kuO;;TEusIoCG?-Pg#Oc((0|qv`p;WJ z|3ypazibKpS1qCcx+V1Aw1j?tOX$CC3H^bV&>w6G{dX;)|Gp*kKeUAY$Cl6^Y6<<} zmeBvy68hh@g#Py}q5nfm=>OOf`aiXV_FAyL37-dqYtzq#>F2^5?Q>zjnP>m&?~7~o zxiH=r^>{|m;~7Ca9}xOIEuqIVf*#KZdORcO@rhZp)$NQol?~8i8FY58WsK@)F9`B2Kyf5ne z`g0=mcwf}xeNm71MLpgZ_5FPuLXY=FJ>D1fcwf}xeNm71MLpgZ^>|;@<9$(&_eDM4 z7xnx5brRa^z~r?+kM~7A-WTD1fcwf}xeNm71MLpgZ^>|;@<9$(& z_eDM47q$E1A`|z{>224huLZB2kv$$i?ym*obK$<81BBoE^mXBl_qwp(%%A@2Ul-P$ z_qw3RGlCw^2zopt==bz-2|b<>^mskCX=9bV)wuJ8gW}Ng}dV_F%>6XyTw1i%^CG>JFp_gw7y+TXq6Uf0glY_HebgzM|Kgx;Vf z^oA{=H);vJaZBh$EulAQ3BBp;n={}3RsGi9B3yr4OX$s7LchHw^gCKYZ{8Ami`cl`WyKY6*RHOXzD_LSNew`ns0T*SCbep(XT!5eRoUfds;$2&@K|O&RqLGtX+SkC4Bwpqb;F7 z))M;TEulZr68e)Zp+D6U`qM3;KhqNWvn`=N*An{kEukN53H?w@=!aWEKhhHV(U#DU zwS<1WCG;0sLVvL(^b;+ipKJ;JR7>cmTS7n668hPe(9gAme!eC23oWr`{iT-hpP|0k z68fc<&@Z=yex)V!t1Y3w+!FdLEup{K68dW`p}*b|`Wr2wzu6LN)Zc0e*T3Bo`a3PL zTK(OYShfCMOX%;ng#JNGtWrBWFJWTS*_rwFujgI6a=7+4^RBsc>D=}A&t3b_@Z%TH zz5Ss#&%SZ-T>sNtUp?3VFxNNC_219+jdT5vwO{Mx*X#eU`kQJ(dly14*%I0>Z|0Xc z*KcVF*OzVyy-Z8!Wm`fo*Aja9me4D-gkG^F^hzzES8fTtN=xWfTSBka5_E#-`W!TZ7rcU zYYF}KmeB8L3B7qs=q*}8Z`l%htCrAPw}jrNCG@r}p|@)Zy?smQceaGyp(XT=EunX6 z3B7Ym=v`Vu@7fZ2x0cYmw}jrKCG?&xq4#PDy?0CKeOf~A+Y)-emeBjRgg&4p^nop* z4{8a0a7*YzT0$S%68f-~(1*8#KB6V`ku9N*Y6*RGOXy=-LLb`_`nZ3H^zd(4TAx z{i&AFpKb~LnU>I>Z3+Fkme8MX3H@M8=!aTDKim@fk(SVpwuFAHCG_Jhp})`)`im{0 zpJ)mFWJ~C$T0%eF68f2z(9gDney%0-^DUuYXo)|sztj?cTEEy5`lXi8FSmq#r6u&M zEup{M68bAGp}*P^`fDwrzupr18!e%~*%E(Nf2$>2|8`60@3h39)ZcB1Kd!&m68ifs zp?}a4_t(z%A2xF=p4;hJe0mn&u;&8dwRWSH&>OdeUepqLla|n%&b~SG?KMkI&*nS! zNDv;MomxWg+!A`1me9Mlgx;+s^zJR8_h<>dXG`e4T0-yL5_+GO(EGN8-mfL}{w<*o zXbF8_OX!1ILLb}``jD2;hqi=1tR?i}EuoKS34LTs=%ZRfAKeoAn3mATwuC;eCG_zv zp-*TDePT=KlUhQb+!Fefme8lRgg&h$^yw|3&u9sKW=rU^T0)=Q68fB$(C4;gOl?daiDX57gJR#Pn=FJ)1w& z^MG(xeYhp`BQ2pHZ3+EYOX$a2LVuwp^cPz~KhYBU$(GPhwS<1UCG;~bp`UFD{aj1v z=UYO*&=S+L`SfhQ^z2h%=CyN~-X^@3F541%xt7q&w}f7yCG?6dp;u}Ny>d(FRa!!? z+7f!Tme8xWgkGa1^qMW9*J=s9c1!4WT0*be5_-Ls(CfE^-k>G)hAp8tY6-n@OXx){ zp*Lv>z3J?mGv7|n=F_j6AJ{X2@ccNaCG_-cK0TXH&gOr+)c;yvH=mx(r(ZXJu-7TV z*>rj~pPtRXGy8Oye%-vsVR|;7eKI)9OwZ=if77zt+-rm9!|uIJ=sj9O@7WT1ua?kz zw}jrOCG@^6q4#SEy?;yS16o2K*b;hrHlKdo{G!>X!_4!+xqNaq|53kg?)i0jACJ&i zw1mF0CG@}4*UcyYhULb79mEIvyU9&0;n#O_OXyo#Lf_gF`nHzPx3`48qb2m6Eursf z34M1<=zCg1KhQ1`G5t3!)32Mq*z|-(gH!ahzn{U`t zf$&kr)Tr&*UhK@rsdnSFGn+fckngyf9|iFe`xOGoBx}=yMVK^zT0(AAb5b_ zP6mQ&a3{gtT|)=~f+PWgTPjeY?%q_O-cqVm_r~4b-5ri>-}iUT^N?TGStsZ0&%59A z?nBn~={5IU^JFF&ThE^%Z6`F(_s{12v-$t*v-!w;Z4u65t7ZwkT9(kOX9>MVme6Zv z3B6X9&}(N2y-t?U>t+ePUY5}7X9>MQme8ZJgx)Yq=#8?3-Z)F>O|pdEG)w5svV`6| zOK3kMCv^YkBwVYz1T z%c1kM|8w%e4~DbNLV1?Z{r~P~|L?H$e{MeG`(bpRC47I3$r5_|ETMPE5_-ohp?As> zdgm;mcgYfZ*DRrT%MyC`ETQ+v5_-=pq4&xXdhaZu_sJ4^-z=f`%MyD3ETIp`68gX_ zp%2Ou`rs^~56Kex&@7=3%M#kT{KIlKUuEw8-(fi`pHIT)_v|d8&&d+{+$^Eb%M$wh zETJ#R68gd{p)bl3`r<61FUb=6+U$S`{l8z{|NG_7=Iemy|NZj*|7PO+gI@}r$BG5= zJmIlp!7QPNWeMFsoA>{I`Few24xO(D=kos9eDDL}`=ft0@1M=bmhX>a^E~0}!{28i z^l@22kIxeNf9Bt5`F+mv6Z3WnZ~x>hp-;&Y`qV6;PsdH=ur`TpRS!w3I+=0ln`Q=8}ikjJ+A|7D~mdlLTGstH1xpw(bPiOppO0_>H@zU$MUT=aIHeBCz zt-j40x~|ptq5Jo1??vA1ja{$T-gjNQsq0#On>Tk|Yd@A=@0K3d+Ryu~J+65j+%}7A zx6k6*9bMPzpYzVHYxQG&SJ$=nOybXSch|KEUf6I?*R}fl-P?6-eEUk-`?{{RKQHr3 z?w`fA2bybFH*d@LcKiMPRIbgLCUkBI9ozQh*7=fau6uKxH`(+&S;Fi4J7w+PE7S9( z3GJVa+kZaR^`DP*{pVvnEWe)6^`DRR!g)EN7s(QOc$UzMW(mDmme7l53B5#?&`V|s zy;PRaOJ@naOqS5gW(mDqme9**3B5v=&?{yMy;7FYBeH}Z$P#+xETKnc3B78T(5qz$ zy?U0=Yh($%W|q)vWeL4@meA{D3B7KX(EeGN(CcRjy+M}Hqq2nFFiYr-vV`6^OXy9q zgx)ku=*_Z(-aJd_EwY5(GE3;KvV`6`OXzK~gx)qw=$z_stS|zbv8m&l388 zETIp~68fMlp%2ay`j9N456u$#uq>eu&l38GETNCg68fktp^wfI`j{-C$7TtAY?jdD zvV=Y^OX%@gLLZ+c^n@&-PskGb#4Mpt$`X2Fme40>34KbI(5GezeOi{#ld^=KoF(*> zETK=&68elRq0h_`dTN%?XJrX}c9ziRWC?w4meA*A34MN+&=+J0ePNc+7i9^3ahA}R zWC?w3c0hzDoAxq1k-jcVJf6NjOXwT2guXFL=$o>HzBx5($8iI{alvN&u0n!LYB}kW(oaLme4O}3H?fz z(643*{aTjLuV)GUMwWOq{brVM{jDsa-_8;br{BpE52fGD68gO?q2JFE52ousTTPV> zPj_8wA5&d^rt6xQ4R{is?Q!kY!Ao;I*W=oW&9&!yTsxz=_Ck+qCp6by>~XDqn=kda zc6{@CFZa0C{(i6YxHh4Ay;pl&YacIP>v3&z^LlU1<(V1(9IM)o_v<}g?@YfSywT&@ zX@i&Mc(cc~_RsQGk8ADk_jZqK?Vt0V9@nNczu&t(uC*Vt_j+7wKkx7NxHi6doBvyn zYZIDl|9jW9jhdI6?vLkYv$*C``!@fNuD7}2|I63=Kf7MftF`a@f9P?o{r&!9k83A2 zzvqAIajkv52j=vPX8wAa*u36u`+1+g$F=tRZ-E}y+PAr2k8ADs%CH{S+SglX7S|T;ajpGa zEYjmz`!;o&D*?e z{&MY$`JP^&s~76(fv#@vk6wIlY`>nhbLO?4X;`uA{i%;9`!DDh&3IgBKdxSSK-abU zJ$zu-wf6Jx^$zN~)_y)*d#Uq|PrKxIA`EHt{MYkw@Nn9BNNw8fxN}QxdT5sL>+_!v ze(tBetDSeZx$fiLe#~8;FHLBFoc42}eIDA+h4!v=-ks*x56cp+`#iOuE4^@<(2HaV zJv>Y3MYDunEKBIcvxHtEOXwxDgkCC3=%urSUM5TEWwV4{E=%a;vxHtDOXwA|gkC92 z=n+{$4`d0wa+c7m3_ja=x^g~|Bl8^L`l?w%ua+hB>RCdsktOt+SwgRsCG^@^La&o0 z^txF>ua_nC`dLD6kR|k}ETNri2)$93&>njTy-AkPn`Q~^eAS+@w8y9RF-m((BJ`G7 zLT{BN^wwEIZ<8hTwpl`NmnHP*ETP9_3B7%m&^u%ay>gdUeA^l@22kIxeN_$;9(WC?vj zme40=34KzQ&=a$SJ~>P1Q?i6UHB0EzvV@+LCG_Mhp{HaCeR`JAXJiR|W|q)XvxGh? zOX#z+ggz%r=yS7#J}*n?^RtA$AWP^AvxL4VOX!QUguWz8=xehBBCI{nV0#(XN?(^H z)=XcYCG-thLf@Dr^i5ep-<&1%Em=a}nkDpYSwi2QCG;IxLf@Gs^j%p(-<>7&Jy}BE znBUwT}nkDpOSwcUaCG-MyLO-7+^b1)+znCTTOIbp{oF()tSwg>>CG=}qLcg9R^cz`X_4J!r z!u7YZgnm0qteSo&ON>mvn`C00p0tIuh|@$sb2b2*dbnx5SBI?b;g()23L zn)SM_dA*kPyROyO+o0>()b^FKqq?q5FW2UqTTHv;>gHwsOzod*tGu4~r+uIFTxmiN z%@W$1>%7UP^P&&V{I1WNCiHw+LeHNi^a5EzFPJ6tuq>e$$`X3vETI?45_))+(2HgX zy;zpei)RVFM3&G?W(mDime5OQ3B62~(932Cyi%V!C_LYB}gW(mDgme3=zgdWHe zdgUykM`j7VYL?KeWeL4{me6Zt3B6{P&}(H0y>^z+>tqSNZkEvNWeL50me3ny2|X%H z=nb=k-Y84xjkAQ_BunT`vxMF(OX$tBgx(@c=q^wcb&&&m?|>@16+zk7Wt{@hqV)%@X>uETKPVIe?Lp;=~+VmAWP^U z+D;AjkMca>`j4}Oz9LKLE3<_DNtVz*%@X=&SwjCjOXy!@3H{3~p?{Sn^slpo{!Nz9 zS7izP+bp54&Jy}}SwjClOXxpj3H`?`q5qU6^q;eYz9viPzhnvh*DRs`mL>GxvxNRf zmeBvq653+i@Kf>bUhls1UGlH&X1YOSvx}FhqJtOFPM$q-KK-b3tUGIy!-WPSfFY0<<)b+lo z>wQtz`=YM*MP2WUy51Lcy)Wu|U)1%!sOx=EPtWIx(DlBk>wQtz`=YM*MSX?0(|K!y zZ?E1LU9a~=UGIy!-WPSfFY0<<)b+lo>wQtz`=YM*MSWHNtc3PB&^{LEdSBG_zNqVc zQP=yTuJ=VUv+)^}eXls1UGlH&X1YOSvx}FhqJtOFPM$q+)pz9ey*E52yj|;k<5p+Ex=z2!b^^BnF89~=G zg05!-UC#)*o)L6CBj|cY(DjU<>lr~$&*zEI^|3(LGlH&X1YOSv`ii_ALf12bu4e>Y z&j`An5p+Ex=z2!b^^BnF89~=Gg05!-UC#)*o)L6CBj|cY(DjU<>ls1UGlH&X1YOSv zx}FhqJtOFPM$q+)pz9ey*E52yX9VrOFxrG^jh=U{o)L8av0%o}(R@5UY>x%~@BJC) z0H5F4;d_6GZ}u>~^K|fu_ub=y$2#~m?XPof&NQLt$`X2LmeBc6$F}FmGlc8&W(hrC zmeBKO3B5p;&JuFM;g|dWRI7{e7vVMYme8AK3B6gC(3@uoy+xMLTV@HpRhH0OX9>McmeAW~3B6sG(C+{389(me6Ns34KnM(C20eeO{K(=Vu9hL6*=LW(j>! zme3bx34KYH(AQ=MMEFV5UWO~v*JTMmKe|3k=o_+xzA;Pao3ezyIZNnUvV^`hOX%CO zguXpX=sU86zB5bcyRwA7J4@($vV^`jOX&NuguXvZ=m)ZdelSbuhq8ozI7{e9vV?v# zOX$b4gnm3r=qIv-elknwr?P~8I!pXG{Y;kd?@*u368gC;p`XtZ`h_f^U(6Evr7WRe z&Jy~SETLb`68g0)p^#JT$%E5*y7fG!_uqHlIoIaQRYK2|CG^lNq5W2!->&Jr=!0{$>+_}wJztj4^JfXY zK$g%8W(hqkOX!8NgkCsH=tZ)G9-bxiqFF*OmL>G!Swb(7CG?V6LNApi^wL>EFOwzo zvROhemnHP_SwgRnCG?6}La&r1^oT5>2eO1-IZNn~SwgRxCG=`pLa&}B^cq=0ubCzE zT3JG`oh9@-SwgRyCG>h(La(1C^afc%kIE8y!z`gU$`X3xETK2a5_;1tp*PDCdh;xy zx5yHD%PgU{$`X3(ETOl_5_;P#p|{HtdUTf1W3q(aK1=8wvV`6-OX!`lgx=Y9YOr_7 z^MvcWW(mDpme9Lr3B5;_(0gVHy;qjdduIu~PnOX8W(mDtmeBiW34K78&*me7Z134KJC&_`woeN>jvM`sCrOqS4NvxGi2OXzV~LLZkU^!O~H zkIxc%LYB}cWC?v@me40<2|Y1O=##UAJ|#=&Q?rCVElcQ0Swc_F5_(FO(5Gh!eMXkh zXJ!dKHB0ETvV=Z6OXzd5gg!S*=<~9KK0iz73$ld1FiYr*vV^`kOXy3oguXUAAi_>f zdl`02Uza6zNMD~N^bJ`;-4sLf@Vx^c`73-q+Lf@Yy^aEK!KbR%-Ls>#UoF()lSwcUWCG=xiLO-4*^b=V^Kba-; zQ&~bkoh7zUKa(X!r=QId`nfEjpU)Edg)E_8%o6&gETLb{68e=apfU(XWy zjVv)H{brVM{jDsa-_8=-rQgXC+os>m68gO?q2JFE+obC`dKKBQ%Pg+#+FYBr{l(dV z=K0-oZT|A|U(G$^yz9k!p|0x#UEO~Fd+|{V41SNdzht)_zxR~p;-|9p_x^R)@2%_a ztq0mSoYfOcUzX?F#}Y5zqxrpuG{0tSvu4lcTKo3fKgXAc3~Vzpub>B_-aEfG?^cKQ zx*mw__4e(w*LA(_d9>FD@;uwuKOxUU?=ob@JworAZa>DlevI|bx$gZQXkLG>=6$;+ z-7=wh=Ff8tKGXhwe`$JBbN!D^pV{;unx5M9pVIm5I(m`+f0H@0g!U?g9-1X|K5ZZT zGn)3t?fj9O>%LCguT|IQOB34Xru}?rpNsbMqZdpQdRUgwJ~!>>OD~)z^deb856=>M z(JY}C%MyC=ETNak5_-ujp_j@Mdg&~om&p=(*({-#%MyC|ETLD(5_-ifp;yWhdPJ7c z16e|^oF()sgU@!JPR=KKWS%2jUo}hU)v|kUOX#(-gkC#K=ykG$UN=kV z^|FLsKTGHhvV__LT{8M^u}32Z;~bSrddL7mL>G&Swe4-CG?hALT{BN^wwEI zZ<8hTwpl`NmnHP*ETP9_3B7%m&^u%ay>gdUeA^l@22kIxeN_$;9(WQkSgp822IewXVf z=H-N6e^Qpv6SKrQLudYHHrG$d%L#pIme8kV2|X!G^#A|4X8vb5AD>I|aUt}z*#Qyy z|KGVY|1(_w|DS8-e}?n_rTcuoP6+)%me4O|3H?%*&@X2R{YsY5uVxATT9(kSX9@jA zmiSrv%`D;iTUkQCoh5#fekV&@nSM7*==ZXOem_fGk@nAU*O-R(v9NzO@1ME<}?2NinI4U`M41J-Yn5SoA=M={j>St_kzcc1@bJR7t9iRSeDQWWeL4-me7l2 z2|YYZ=tZ-HUMx%K#j}K7B1`BcvxHtMOX#JugkB~~=w-8nUM@@M<+FrdAxr4~*}VUA z^KA#e96DbQ+vRz}*T?8Ap~qwiy?vI@J7fvHW0uhUvw8pL=Eo0yIdr}r{G6!$S<(OK z+1%IDzvj=)_s!XVa{he8Z}a<2$r675)3by=BTMKrv&2>DsafJTr6(`;v8}F2pOu#r zUjOVYq0h+@`rIs`>)BlQ&*n4!?vu0f?fEzn`i?B2@5~bVt}LPN&Jy~bETQks68gR@ zq3_QU`hhH=AIuW^p)8>v&Jy~OETJFG68f<$p&!o@`iU%|pUe{asVt$N&JsUQKa(Z= zT>057p`XhV`uQxOU&s>r#Vn!wKQ|xzSoHr+OaE-XQoa-jkL@F}gdWHedgUykR~dY^ z^R)kS^Zw7xcg$A;;p?M+Ht(O!|FzEMzP|qd@^ka{EMCv*`{kbn`e*YQ|Gflf!JG29 zLFk*aguW$9=v%XdzAa1W{@J{LHh+Ea%b|1je_#W3zXkPdD9@m;nbIj`6gju}a30>Ff=i$Vz zYxU!KQrETi=Uo0g6T7a}w|R2cwfgZqrR!Szvn_A))UIpo&&^yrt?OESo0FPr!~McA zt?pC$D4ETMmpCG-!og#J;M&_B)+ z`id-}ugntqCs{)OG)w58WeNTBETMmqCG;<|g#J~Q(7(%x>QcV>(~`Gv5~<&-X&TFY0<<)b)&@oev0oX_nCSjG*foLDw^au4e>Y z&j`An5p+Ex=z2!b^^BnF89~>_0$m>qbiFU?dSBG_zNqVcQP=yTuJ=VUv+)^}eXwQtz`=YM*MP2WUy51MH`{IBJ(;7YRTK}ls1UGlH&X1YOSvx}Fhq zJtOFPM$q+)pz9ey*E52yj|;k<5p+Ex=z2!b^^BnF89~=Gg05!-UC#)*o)L6CBj|cY z(DjU<>lr~$&*zEI^|3(LGlH&X1YOSv`ii_ALf12bu4e>Y&j`An5p+Ex=z2!b^^BnF z89~=Gg05!-UC#)*o)L6CBj|cY(DjU<>ls1UGlH&X1YOSvx}FhqJtOFPM$q+)pz9ey z*E52yX9VrOFxrG^jh=U{o)L8av0%o}(R@5UY>x%~@BJC)0H5F4;d_6GZ}u>~^K|fu z_ub=y$2#~m?XPof&NQLt$`X2LmeBc6$F}FmGlc8&W(hrCmeBKO3B5p;&JuFM; zg|dWRI7{e7vVLh4Jt|A+4YP#aC`;&#vxMFxOXy9rgx)Mm=*_c) z-XcrrEwhB)Dog0CvxMFzOXzL0gx)SoX!rm244}uP3B7%m&^u%ay>gdUeA^l@22kIxeN z_$;9(WC?vjme40=34KzQ&=a$SJ~>P1Q?i6UHB0EzvV@+LCG_Mhp{HaCeR`JAXJiR| zW|q)XvxGh?OX#z+ggz%r=yS7#J}*n?^RtA$AWP^AvxL4VOX!QUguWz8=xehBB0Sl& zm*I)@by>pCkFL)W`i3l_Z_E<NL;`i?B2@5~bVt}LPN z&Jy~bETQks68gR@q3_QU`hhH=AIuW^p)8>v&Jy~OETJFG68f<$p&z%M8tfJUc53tIoHmPV zAL((e{rx`Lh(?Z@TgUB91e?fZUd zk88*I1>v$D*V?a%PxQFf-futIb#3zCo9jINRM)jr%C%2-U2~~@-OGDiYd@Z!>2a++ zE1cZ*HQ0V$z3tETc)bA^!soiKwI55@KHqiC%LY6NU+8hI{kr~Qk8ADs!hF; z=lF7uYXi--ugv1wS7&kUYdx;D@8j2dTx);7Z_MJ_H+x)b-{!Y^Tx;Lvw|iV`U++6T zuC=ds-B2Hd8Q&|aeZB9_;`P4Q<68T%{(g^Z?dNZLk8ADM;17CSYwuOVdR%Ki??3Fi zHr+eX@WrmL-S&II<(E1)yBY6$`?%xUISYC^<63)8_R@2^t_^r$!+me8Nb68e)_LVqet=uc+}eR-D9pUD#XvspraE=%ao zX9?|NIL3r&jh=UHcKcX-A@4Wg{r_T?&|k_D`pa2De-a^emx&kR|jFvxNRpme4=W68efPp|8vm`X^aJ z|1?YJpJfUC^DLo%ktOsmvxNRtme9Y>68bk;LSL07^l!6-zB)_j-(?B?`z)dVkR|jV zvxNRrme7CB68f4fq5qO4^k1`t{#%yNf6o&7A6Y{GGfQZX1*1&xeL%R@|6b^SFU)S= z3u8K8{WISe*Ua}qy)Wu|M$q+)pq&p0eQB1^mt_fE&j`An5p+Ex=z2!b^^BnF89~=G zg05!-UC#)*J{IWuSfJ~DQP=yTuJ=VUv+)^}eXUv+)^}eXD$lD?Gm03d9`=YM*MP2WUy51Lcy)Wu|U)1%! zsOx=E*ZZQrDt|6QdmLyV3v|6N>Uv+)^}eXwQtz`=YM*MeV*gV8XOU&%4%tESPad_VM_zKNi%lh3WYkAiVGW$A#JNabZm7vp@6W z!q?}X`EfzlGlH&X1YOSv`qI2zLSL38bUh>JdPdOojG*foLDw^au4e>Y&j`An5p+Ex z=z2!b^^BnFJdPdOG z^LZh3eJs%RjG*foLDw^az9Mgj&{t*&UC#)*o)L6CBj|cY(DjU<>ls1UGlH&X1YOSv zx}FhqJtOFPM$q+)pz9ey*E52yX9Qi(2)dpTbUh>JdPdOojG*foLDw^au4e>Y&j{Lm zVYCU;8a?k?JtOG;W5JA{qxpDz*d7b|-}^Jp0Y1O8!}tCW-|S&}=jq@R@4Lr^_S%Hh z*@Ktn_{lu~*K6aNYd`I}HqgA>{H*Jmx7qUZu4@DBh3qf7uGQE3WpizK^XsQ&iA{3t zBUwV%$1h#aGP<5+^rd+{LSL38bUn-HdX~}kETijLM%S~9u4frt&oa86Wpq8uXrHq& zCQNJeylb=D$D*F=y#MuFr|Y>+*T-^Q&vm+<>vTQW>3Xiy^<1axxlY$}ov!CPUC(v8 zp6hfy*Xeq$)6?^LC-e`pgs$g0UC(v;$N6=Hz9LKLE3<^I=Q>@_b-JGGbUoMUdal#; zT&L@~PSvVl=*Y#Yd>$y(XbDgf|I$h6ox}NKFJ=f`auG95gr|Y>+ z*K?h&=Q>@_b=qUWC=>eM3;pi}-w*J$F*|-QjOo1hGv60C&&Q(P7j-=&=z2!b^^Bk| z&D$mPWm!VkGlH&X1YOSvx}FhqJtOFPM$q+)pz9ey*E52yj|IBk7j?Za>Uv+)^}eX< zeNor@qOSKvUGIy!-WPSfFY0<<)b+lo>wQtz`=YM*MLj*A7ed$jqOSKvUGIy!-WT;1 zc{_x@GE3-sU)1%!sOx=E*ZZQb_eEXri@M$yb-gd@dSBEY2inI1T^|c{y)Wu|U)1%! zsOx=E*ZZQb_eEXri@M$yb-gd@dSBG_zNqVcQP=yTuJ=XlzBpjQv_{Xn)_*LR@$W=@ zJU;A?1@&uTdcFn-?|c7oVaCrwd^~1{$AvMS&;HDh3tM$PF6eqj(DjU<>lr~`nzu{n z%d&*7X9Qi(2)dpTbUh>JdPdOojG*foLDw^au4e>Y&j`An5p;cA(DjU<>ls1UGlH&X z1YOSvx}FhqJtOFPM$q+)pz9ey*E52yX9Qi(2zq)xFNCg-1-hOQbUh>JdPdM!ls1UGlH&X1YI8sbUh>JdPdOojG*fo zLDw^au4e>Y&j`An5p+Ex=z2!b^^BnF89~=Gf_7gRZNjuh&%0L72)h4RF!*@TetbS` zj|Khj{Tb%~U+=TS_x=!H%wc-x>EM(0alvC9{F?UHxi)8-&~s%8Jv2+`xwC}Me?Iu% z5Yh9d2|Zty(DP>ry+D@G3uXyDEKBHxvV>kZOXx+igdUzH^rBfpFP0_r;#op3ktOt! zSwb(BCG^r+LNAjg^s-q(FPA0s@>xQ!kR|krSwgRrCG?0ap$D>rUO7wXRR*8!JpEOA zWS%2jUo}hU)v|kUOX#(-gkC#K=ykG$UN=kV^|FLsKTGHhvV__ zLT{8M^u}32Z;~bSrddL7mL>G&Swe4-CG?hALT_a|HP~C{dBXK=vV`6?OX%&ggm(XL z&j5N%n$X*43B5y>&^u-cy;GLZJ7)>KOP0{PW(mDpme9Lr3B5;_(0gVHy;qjdduIu~ zPnOX8W(mDtmeBiW34K78&*me7Z134KJC&_`woeN>jv zM`sCrOqS4NvxGi2OXzV~LLZkU^!O~HkIxc%LYB}cWC?v@me40<2|Y1O=##UAJ|#=& zQ?rCVElcQ0Swc_F5_(FO(5Gh!eMXkhXJ!dKHB0ETvV=Z6OXzd5gg!S*=<~9KK0iz7 z3$ld1FiYr*vV^`kOXy3oguXUAAi|bSdl|M!Uza8PyyW^Up>N0%`o=7wZ^{z-<}9Ia z$rAe3ETM1968iQmq3_5N`pzt&@5&PT?ku72$rAeBETQkq68ioup&!T+`oS!rAIcK? z;Vhva$rAd}ETJFE68iBhp`XYS`pGPzpUM*Y=`68X`k5@@-=RI5CG>MyLO-7+^b1)+ zznCTTOIbp{oF()tSwg>>CG=}qLcg9R^cz`X)AXBJ!u7YZgnm0qY@B{4OKg;WH%sXE zvV?v=OKh0-&(uRqm^3(@r@wBl4R2mHElZ5fwU1;8{n0F;Kb9r*$FqdKG)w5qvV{Ib zme8Nf68ckFLVr3-=*zQ&{!Et8pUo2bb6G-vK1*mHr!gi>YxKNpv)jkw3wgf@@BbIG zg#J>N&|l6H`YTyNe>F?!uVo4S^(>*kktOsuvxNRumeAkM68bw?LVq_)=Yw?(xI?}d>U~kyGlH&X1nqo4 z=u5MNzAQ`VdPdOojG*foLDw^au4e>Y&j`An5p+Ex=z2!b^|3(L#{ymNi@M$yb-gd@ zdSBG_zNqVcQP=yTuJ=VUv+)^}eXwQtz`=YM*MP2WUy51M{Rrzxf+T%d`SfJ~DQP=yT zuJ=VUv+)^}eXJdPdOojG*foLDw^au4e>Y&j`An5%lzYUI<+u3v@jr=z2!b^^Bme$lD?Gm03d9 zGlH&X1YOSvx}FhqJtOFPM$q+)pz9ey*E52yX9Qi(2)dpTbUh>JdPdOojG*foLDw^a zu4e>Y&j`An5p+Ex=z2!b^^BnF89~=Gf_7gRZNjuh&%0L72)h4RFyrTFJ{}*o$AbR% z{)}^g&+qK;y+6b^dzjvNI{3uMXme9kqgkCgD=*6;xUOY?aC9;HGGE3;C zvV>kbOXy{?gkCmF=;gA6UOr3c6|#h0F-z!`vVh(La(1C z^afc%kIE8y!z`gU$`X3xETK2a5_;1tp*PDCdh;xyx5yHD%PgU{$`X3(ETOl_5_;P# zp|{Ht+IgTo6X-E%LT{fX^bT1<@0ca@PFX_lY&$jByX1Mo^OX#Dsggz!q=&@NsADbogxGbTM%MyBgme9v%2|Xc8=o7MpJ~2z^ld^=K zm?iYdSwf$ZCG@FTLZ6l;^rS4ICua#gB}?ejvxGh)OXxGRgr1ru^jTR#pPeQ2IaxxV zn4COJEpJ85`La^eU{KSWC?v^ zme4n434L>x(6?j>eQTD`w`B=^dzR35WC?v|me6-)34M2#(D!5seQ%b~_hkuvf0ob> zWC{IXme3Dn3H@-E(2ryZ{b-iZk7Wt{c$UylWC{Ibme5aS3H@}I*gpMCmhkUXpUo2b zxh$cd&l38DETLb_68fbqp2*dWkHdm&_7+sVt$F z&Juc=ETNap5_-8Tp_k7RdW9^ZSIiQ6r7WRGWC=ZxCG^T!LXXT6detnUSIZK5^(>*+ z$P#+ZETPxR5_;_{q1VY0dfhCc*UJ)m{Vbt3$P#)~me3n!3B6I4&>Lq7y-AkPn`Q~U zS(eb7X9>MUme5;f3B6U8&|7B-y-k+T+hz&9U6#nOQoETPZN68fAhq0h|{`n)Wm&(9M2f-Ipg%o6&dETJ#X68e%Xp|8yj zh;U8QUWPxXugemDN?)HP^bJ`;-4sLf@Vx^c`73-q+Lf@Yy^aEK!KbR%-Ls>#UoF()lSwcUWCG=xiLO-4*^b=V^ zKba-;Q&~bkohAO5ekM!&KK*Q#(9dNF{d|_tFJuY*VwTV^WeNRqme8+c3H@r8(640) z{d$(rZ)Ax-q~FXEuD_Kf^xIkDcj4sYSVv8d*9pl|JP|B)Ar-4y$Yd+W(n;>*!eIv?T_2} zBR4&7mhkK6%M#khuKoCFAG`MBrx#2UdRUgwKDO=0S1+6<^deb856=>M(JY}C%MyC= zETNak5_-ujp_j@Mdg&~om&p=(*({-#%MyC|ETLD(5_-ifp;yWhdPJ7c16e|^oF(+g zETLD;5_+{Pp;ylmdW|fh*US=ntt_F}&Juc^ETPxU5_-KXq1Vq6dV?&XM`a1UVV2Mv zWeL4;me8AI3B75S(3@omy?K_CH(r+vxGh)ON^L%=6~L~D%Ypx<-~7FPhRX}TV0brD=#O! zy|c4~J||1)bF+j#FH7k2vxL4NOXv%;guW(2z9dWNYqJ9){Hkd$!!OgNF+`nD{gZ_g6?jx3??%o6&pETQkt68fGjq3_KS z`o1io@6QtYfh?gP%o6&cETJFH68e!Wp&!i>`mrpbAI}o{i7cU?%o6&kETNyy5q^fFmOFPkOwa#=zzpC$APSwgRvCG<*JLXXH2 zdLT>am9vB%nI-h9SwgRtCG_f9La&h}^qN^huazbA+F3%clO^=JSwgRuCG`4PLT`{I z^r$SMH_Q@xqb#8}&Juc)ETK2e5_+>Np*PPGdW$Tfx6BfHt1O|n&Juc?ETOl}5_-EV zp+{#4Jtj-&?X!g5Axr2TvxMF$OX!`mgx)1f=v}jf-YrY$-Lr(=BTMK#vxMF&OX$6` zgx)7h=zX(<-Y-k&{j-EVAWP^2vxGh zOX#Dsggz!q=&@NsADbogxGbTM%MyBgme9v%2|Xc8{7cW`&dMj}?GWDnDOp0FnkDpU zSwc_B693Y(xVLv^-VV_}i}%ms_vhyULs5AC9{NHDog04vxHtIOXy{@gkCO7=;gD7ULi~9 z6|;n1DNE=PSwatF3B7Wb&?B>iUNuYT)v|={pT+xU@%~wS#`nXad6w|~aafknhi3_W zM3&G;W(j>%me5CM34KhK&||ZNJ~m6}aaltDe|Z)^edzz%KNt7$_}4s(Ps#g#Za(gW z-|xIEF*1FAmT>)oETJ#V68fSnp)bx7`jRZ6ugwmK(Equ3|L5Y*=Iem)bL8i;gnm9t z=ohktelbhvm$HO@IZNnQvV?v$OX%0Kgnm6s=r^*&ztzvhhYbGqAN(`B{Kq_=&iLQV zRL6S~{@LT&)WJ)0{JXAe?avil{`Xzi+ShUI+OBI})?UA^$F-CEf^hvTuH7(;Yd3aX zn>_gDI!|xvx^_ajc5`#BGkJM?*7r+}$_v}Ge%tL?zU=|eH2C$aw2{3v?eVPTmdstG<@=I>-x>i32m*+9NeqK+N4R>_C-avaHduP|RDdpN-UDwVi*Y2Lh zwR>i9?cP~jyRW&{nfzY$b8uo_SU(54eh&0N`Sa=eId~vlKL_`x>&I)W^12UpU2D%F ze#t{!*9OY#J=}F|e7W{W*R}ewd~_Dq9_zYRKj)8kUGrxlp6I%E;^2!rPoM0%R=>uc z>bf?)yxy~Oiy2?Hs{Pu1y~nkQenEJ;>uq}3fG6RZ9@pAG&$B(QwI8?VdR#ke@J-I~ ze2;7G+kBzNwf1}E#pc?b&D$~FzY&*SK6-)9_t)UxtMw*VuOZGmW;}pJ}w+ z{!F9oW%9agd)YkCwwKHEY#LXRLvmfSy}m}d?h8+|y}o9i zXWMJ#dA7ZFo^SbY`SWYG*Vkz;X8&uhYqrH!5tRv z5Zoa^aCdh?2#^Gl;Fd}Xb$45plzJ)D^_&2uEmUxKcYXcc@4N5MTJs^_%sqGh*kg}z z&LHa=y>icM{jwxW=I?=0%65H|a{WJZU9(-^v|Rs>T-R*ZH!IiwJ=Znc_0f6E+Wwly z&+otG_ivuZY}dEQW47KhkJ);wJZ9^y^O&u-$z!(OHjmkQyF6y=?emze-NV>=$2?~1 zo${Eich2L+zvj=Q*{<)B$85c89yk6azhARmACt#yy;~l$_3nAh)_deJTaV3Sw%#+3 z*?O-$X6wE4n63B8W47KmkJ);^JZ9_t^O&s<$YZuXFpt^#pgd;lgY%fJ56NS;J~WTn z`mj7^>%;Syt&hlKwmveC+4`tFW@~@{hOLjuW41mvkJV)mbMU9(+3C660#<@zZJw(FJv-6m(&&gx9J~xlqdSV{4^`tyz>&bb{*5~ChTc4lDY&|89 z*?MXov-Pw*X6p;`n5{3&W4691kJUryhc$J8&SZ^~n~zB!NC z`j2_c*0pSz9t?$ZXw!S-$+4`P5X6t+Nn62;2W469OkJEIFH%- zkvwMWNAsAiAIoFv7xVReJddfLPd|~zZ2e>&v-MMX%+^omFA$87y<9<%jxdCb<& z=P_HqkjHHOVji>gOL@%JFXu7!3;Fn8$z!(bujVmZzm~_;&*k^Op2yVBrr*e8wth2@ z+4`+KrhX>Bzx%J{Jolp4U&~Dyy!5VT7xG}vzXojAp6hXKT<7(k?{TfYzZd3m?ZvKZ zTXx=_@z+fCW7$@A=)SgHyO8a^=2H7uUg~;(uC+hk%RR2069-@IUH{bM+UcEZ|J>u+)Xuek>2a;S z-oN&^c4p`G{;kKg_Gkb19@pCY`;Q*iCUoBBf6nFFdJFhXbAIpkbzbj(^msk@Qv6qs zYwhds-#xDN4ZK!+4*bZ!Xv7?{V$4xxCHxIqUUjXy2a;^mv=o zI&X8a9)G^}dY3KeH_iU@weK-6eRlZ&{@V8!o%34zT=sfBFYEJ5Sh(xsoRpWf=a)r# zTswE*wet5kx~`qkdAS+Ub|7g|&hPK*>=Jq5z^A6u zKg{QMMCY;5BfGkbz{a>OA(pnSOU36Z$<_Lccdl==Ws_JtIr#_h$)x zb(YZAWC{I&ETKP`CG>}~g#K`r&>zVX`lDGwe=JLApNlajOz-HiYxCOY;^X;q6aM_4 z$P)UKSwep*OXyE$3H_NYp+B1?^yjjK{(P3uU&s>ri&;W{DNE=tX9@k4ETO-eCG^*_ zgznE0`s-Oje68bw?LVq_)=wQtz`=YM*MP2WUy51Lcy)Wu|U)1%!sOx=E*ZZQb_eEXri@HBw7ed$jqOSKvUGIy! z-WT=Eyd6S+J4@(#U)1%!sOx=E*ZZQb_eEXri@M$yb-gd@dSBGn<>Ml>=YjTIpzD26 z*ZZQb_eEXri@M$yb-gd@dSBG_zNqVcQP=yTuJ=Vd*PS;NLY^pBHpJBj|cY(DjUJdPdOojG*foLDw^au4e>Y&j`9cFX(zk(DjU<>ls1U zGlH&X1YOSvx}FhqJtOFPM$q+)pz9ey*E52yX9V4!uM45;bAhgB1YOSvx}FjA%)A{! ze>+R)dPdOojG*foLDw^au4e>Y&j`An5p+Ex=z2!b^^BnF89~=Gg05!-UC#)*o)L6C zBj|cY(DjU<>ls1UGlH&X1YOSvx}FhqJtOFPM$qmHyP7b)qsOidJr~SABlvv0)1M25 ze)ead1AM*bg`fRle%M3*z^{Y9wC4rSb?|%I?{jUwG@<9u5_*9wp%=^&I{z84^}=~T zxV}i1(8IHY9+4&VqFF*OmL>GaETI?A5_*X&p+{v2y=0cqOJxbYbe7P|WC^`&me9*( z3B7!l&?{sKy<(QoD`g42a+c7mWC^`$me8wZ3B7uj(0y4#uaPD6nu8AqeqAbGxV7?# zaDDA8q1VY0dfhCc*UJ)m{Vbt3$P#+PETK2b5_;n-p*P7AdebbSH_H-wbe7PYX9>MU zme5;f3B6U8&|7B-y-k+T+hz&9U6#<>X9>MSme4z93B6O6&^u=dy-SwR?*Hu!pvR;M zy<3*hyJrc#N0!iIvxMF=OX$6_gx)(#=zX$;-Zx9={j!AKKTGHXvV=Y`OX!2Lgg!V+ z=tHuEJ~T_{!?J`vJWJ>!vV=Y|OX#Drgg!b;=wq^kJ~m6}E#aTjMk|p$|SwdfyCG_Q4 zLSK<3^p#mcUzH{FtZbhML;rsRbN=s&{r$2#^LZikU0Fikoh9@=Swi2NCG>q+Lf@Yy z^aEK!KbR%-Ls>#UoF()lSwcUWCG=xiV(9;GVD>*V>ht0MPvDI#q2J6B`mHR{{{IGs z4d%Q0J8SmN@AcQ)?OJhfLJ z>f5}f>*H*%zZ)=%)y5PzwVXm?fKpd_D;9wOWpm?kL{W3^*OTr{~hGT z1Jt*@a@Xgi-FN)PRl2U#ufeKa*ZK$E^58H2x6{6luJ63e`)%)gy}aIhX+qDRCG-MW zLNAylv^PKSE;^kTz5Ul`t}l`%w69(JdTL**_Vv??rU|`Rme9U-?dz!*PZN5HETKnb z3B6>N&`V_ry>yn)%VY_?Y?jc=WeL4}me4C?3B6*L&?{vLy>gb&t7HkiYL?KeWeL4{ zme74!La&h}^jcX$ubn0II$1)mnzA{VbtFnZimF*K@-A;QM)=A%(CDu;glqK}dSwjCYOXyp& zg#J^O(0|Sn`Y%~R-snvAwqe(`_OtG{Y}9qFJx{x~ao4pS zys*P2bGf$ZT&`_4musWDuJyNfmL0#4=>69aJLL7|%T+?ppC$AHSwb(ECA2p`@Gd%? z7rp)05UwwhCiL(up+{s1y=a!ui)9HtGE3;ivxHtEOXyKqLNA#m^io+uFP$ayGFd_| znLn6y-}9X8)pf*me7Z134KJC&_`woeN>jvM`sCrOqS5cW(j>-me8kX34KPE(Bra%J~K<`@mWHjl_m6qETPZN68fAhq0h|{dSaH)ld^=K zoF(*mSwf$mCG?amp{HgEJuOS<3$ld1FiYr*vV^`kOXy3oguXOO=*zN%zC26lE3$;X zGE3;IvV@+M?Gs`9PJ0=)OW&9!woTuZCG^c%LjN&K=v%Uc{!^CFf6fy6FIht0nkDpY zSwjCcOX$C43H|phq5qL3^zB(f-;pKsomoQPl_m7uSwi2FCG@>nLf@Aq^!-^vKaeH# zgIPj9lqK}TSwcUOCG?|NLO+%zwn;ypCALaGktOt#SwcURCG^u-LO+uw^s`w)KbIx+ z^I1Z_kR|ksSwg>*CG^W#V(auPS;F;KvxI&vOKh2bJxgqnej`ihH?xF(D@$yiu73^D zCp+xeb!|$!klm^4+M4Cs&Ry50m210nU8`?%*RE^h%Il5kx>n!jZe7<-FR!wc8-Q(KS&b57dT$|Xrwr`JX6FS%S>v65G zb8Y_~*V^kH(BoSBSPtxQZBplL9@OJn`x-v8kl!@>zXMe5&v0;$*E`p52#55z*6!bj z_P92A@YUY6`T`!z`Tp9+d03CvYj5-L9@pCa?uZ`O+Up(J<68TC9@XPo`#6v8ajm`0 zV|rX`uXk*ZYwi6V*W+4yz2kdaYoF^AdR%LtvlDw$>J;?dyI<*W2{@YJY}tJ+8I$*O^_{>c=v^>stF7c)w?LT|2`IJ51=hR-XsY z?z-0BUMc&-&d>Y6mp|Q}Km3*-4V!)D)8Adtud{zv`(=mr@BhfF3{d|}=XBoxz`N>P z-?+0Ab8UF%u~{lTqN_)CbzfJ{Dt&I(pS69h{FX^w*V@<0waH!A+F97O^SZ9JbF*vb zcU`NWzbReU>ie78b*3Swhdq68imFLSLOF^fgℜ`Zf4`vDdp)8?4oF()}vV{I< zme3!|6599p7!#&<^w_m|?Q`+*{J9B#{!e5H{mCq$Kb0l)r?Z6qOqS4}%@X=^Sweq4 zOXx3T3H`+^p}&+R^p~@Q{z{h6U(FKwYgt0~X9@lFETO-VCG5}E=%a&X9@j>ETKIYj5fi~0pZ%v&%)5p z!o2pgFlOMpf6n{j+WA?i_eEXL2)dpTwDSR>XJiRo&j|YJyqwVWjG*foLDw^au4e>Y z&j`An5p+Ex=z2!b^|?UT=K@{ti@M$yb-gd@dSBG_zNqVcQP=yTuJ=VUv+)^}eXpUv+)^}eX3)1M3K_dgIBj|cY(DjUls1UGlH&X1YOSvx}FhqJtOFP zM$q+)pz9ey*E52y&kMSq5p+Ex=z2!b^^BnF89~=Gg05!-UC#)*o)L6CBj|cY(DjU< z>ls1!=j%Y|`dpyv89~=Gg05!-Ju`2I(DjUls1U zGlH&X1YOSvx}FhqJtOFPM$q+)pz9ey*E52yX9Qi(2)dpTbUh>JdPdOojG*foLDw^a zu4e@8zObta(>r?XT0JA^q343xe;>`~$xa;-mwM)9LT~=Q2($2L-%FE8{Jiap5hL@N7XQM}S^~kR7 z>*@)mFU#vsN>A*(?27cPa{cmj{j*Og*RJZi)>p1gpUbs(&*j>C=5p=5ook~y?`L|J z@Xx2bJ4@*L{H*I)P1m!Uo{`rhbUmx-tMhU~*Rz_gXEj~VYPz1)bUmx-dREi*tfuQ( zP5ZixF=2W~k6oMBJ{R@e>d#-#t-79Db-h>VdT!PA+^XxjRo8Q?uIE-=&#k(iTXj9R z>UwU~_1voKxmDM5tM1R&ozP#;68aliLf3Pv{#IU2=$TnU*K@1BHZLc1J-6z5Zq@bN zs_VH`*K@0`=T=?Mt-79Dbv?K0dT!PA-lglgRo8Q?uIE-=&#k(iTXj9R>UwU~_1voK zxmDM5tFGr(UC*t$o?Er&g3%@n{VWXqEckhV?~Qr!XJO31pMK8!qJM_B-WPQ}Bj|cY z(DjUls1UGlH&X1YOSvx}FhqJtOFPM$q+)pzCviuJ=VUv+)^}eXwQu8=j%Y|dSBG_zNqVc zQP=yTo|(5p=z3q&*XHGfuJ=VUv+)^}eXAM?WV!kB@t{+!PX{&!vVc|q4Rg05!- zUC#)5M&2%=>lr~`otG23o)L6CBj|cY(DjU<>ls1UGlH&X1YOSvx}FhqJtOG)yrAnD zLDw^au4e>Y&j`An5p+Ex=z2!b^^BnF89~=Gg05!-UC#)*o)L6^z7B-0&jq@k5p+Ex z=z2!bGxK%`UC#*m+Ps|5^^BnF89~=Gg05!-UC#)*o)L6CBj|cY(DjU<>ls1U=K@{N z2)dpTbUh>JdPdOojG*foLDw^au4e>Y&j`An5p+Ex=z2!b^^Bn17j`vadPk348+tC7 zeMa#4c&9%X4E^lSJ_q=E&kH~M!~C#^{()Zye`(JPp6lTEwBP62d}%_@pC$AHSwb(E zCG5_-ujp_j@Mdg&~o zm&p=(*({-#%MyC|ETLD(5_-ifp;yWhdgUykSIH83)hwY`%MyC^ETQ|dgkB>{=rso) z4*dGQ^jdjDxW0Cl(CcIgy>6D!>tzYOewNT1WC^`tme3nz3B7Ta(3@lly=j)vn`H?- zI!oxyvxMFvOXw}Lgx)Gk=&iGa-X=@vZL@^lE=%a`vxMHk4%A@pn8$?cJ7o#IbC%G% zWC`v5-_8JfOq$TUWeL4|me6}-2|YGT=smN9-YZM!y|aYgCrjvkvxMF+OX&Txggzik z=mWEaJ}67*gR_J_BunT+vxGh@OX$P1ggzom=p(a)J}OJ-qqBrQCQImJvxGh_OX%aX zggzll=o7PqJ}FD+le2_AB}?d2vxGh^OX$pP41}_$;B%$`X1)me6Ns z34KnM(C20eJuyq@Nm)Wq&Jy~(ETPZO5_(FO&{MO7o|Yx_1zAF0m?iW@Swdf&CG;g( zLSLFC^krE>U!EoO6D_xwfZ(c)OD>rcX_=JcU^1GgRXs~b8R6J zh7ZobuVZsb|b7pSzFc^=FmqAMN@$Tx$2IkN3FNKF&|{xYnM#KiT72yRUt! z$F+73`*e?M?SAr^9@pCY`|Mn`roG-J17Dcg-+w-5eSQgF z=5sau^?P9l|DKfkdaBPa;m2Lq>gVLTu50ak-Ru2iF4wN_y4JqlUhkr=YxQ$^!(6Vl zfA7dJZzQxctH(d>dcE;p*x_eg*XsS`=Uvz8=l-=uy2Ozd3$UAn&S-=^#LVf|kFeb=@A&f75lS?l_GzGvF|{X^F^zq5V6 z{XbpT>et}^?YdS!XS3#V?Z&xWyJ;@hZth&`^)>eQvFhht*U!7|^9Spo+h0?8(Jkrv zdH!R%e!l9T`%hihcJNy}{JHB|d+vAnFJ0H#=h(GdyROx*>up`v+I`*Y{k7{_{W$;D zb*+9ZfA6~1{{Dft`H!w^{a)B%Z0FkbogS7Y9>}%%vV@*LOK7*bfm>at7t9i_``Wgz zug-to{(H-Mku;(G?6sdi?R%hoALvEXgkCI5Xx}64`$R9ECiD_nLXXN4ddV!Im&y`) z=`5j_$r5_mETNan5_J4@(wvV>kYOX&5ogkC>O=nb-j-Y`q(jk1K^I7{eFvV`6=OX$tAgdUwG^yXPY zZ;>VRmRUk?l_m7nSwe4dgm;mcgYfZ*DRsO zWC^`nme9Lr3B5;_&||ZN-ZM++y|RSfJ4@(&vV`6@OX&Twgx)_(=mWBZJ}^t@gR+D^ zI7{e5vV=Y~OX$P0gg!h==p(X(J~B(_qq2lPI!owdvV=Z1OX%aWgg!n?=o7MpJ~2z^ zld^<9IZNnMvV=Z0OX$=rgi}9+xHbnOQ=Q&l38qETJc434L~!(C1_ceQuV} z6SIV#lqK}!ETPZK68iisp{HaCJvB?{X<0&FkR|kmSwdfwCG^EvLSK?4^rcxsUzR2G zNI-`j1&c-;yQtpR$DhbC%G5 z$rAe3ETM1968f)MLjNsG=)Y$P{f{i6Z_g6?jx3??%o6&pETQkt68fGjq3^W=HQ4v% zG2#0CSwcUMCG>+?LO+xx^ut*~KawT%qgg^fmL={+>w4GOXxSV zgnlba+@7v~4WZ8UOL(xywf4OFP>*X924C%65BIn>xpVE29@pCYd$h;3NuAewtjD$X z{vPjf?VQf*y|kbQbN)T6_BNmB@p`9qUhl~s*V>=`sUFwbpW*2q*T#3=X0N}O*ZvIG z^!Pa2+wArC$WHIP&1dHF{+{h|tvyfodYcnFZ}Yhxuh;$zb9=pB*V=uf*R?Y`@9+7p z_vdrczF%JGacyen^MM3&HtW(mDmme3=!gkC&L z=q0j*9+f5Zl37A8l_m7jSwb(9CG@gcLNAvk^zvCkuaG75idjOhlqK}aSwgRpCG@IU zLa&x3^y*nc_hpIyNUxD4Twg0o=(V$iUMEZFb+d$CFH7k4vxMFtOXv-=gx)Ai=#8_4 z-Xu%tO|yjFEKBIoSwe4~CG-|qLT{NR^j29yZ=EIdHd#V%nMemeBiV3B6yI z(EDczeL$Ab2WAO6@~IzBxFEf68gz3p`XeU`spm8pUD#X*({--%M$weETLb>68gm~p+C0`MTa_eZ5OMuXlatZ4Jv3i{t_kQL6*=PW(mDfme3n#3B5^{(3@rny;+vfqqBtGJWJ>;vV`6;OX#h# zgx)$!=xwru-Zo3zme2=h34KVG(1&IT zeOQ*zhi3_WM3&G;W(j>%me5CM34KhK(8p#8eO#8%$7czBLYB}cW(j>#me40>34KbI z(5GezeOi{#r)LR$MwZayvV=Y}OX%@gLZ6i-^n@&-&(0G1oGhWw%@TTIme7;3gr1xw z^m$oApPwc4lq{jAW(hqlOXv%-guXCK=!>$1zBo(hOR|K%G)w5qvV^`oOXw@IguXIM z=&Q1Xo|WwrVbM-|8AhaU%o4-XH)RQZbC%G5%o6&RETR9DCG?-Og#JsG(6?p@eOs2$ zf6Ws5Z&^bBJxl0+WC?wHme6-(34Ldl(064CeRr17_hbouZdy?!h zvg=xZyO6zfVKMtmKQ1-Fi^AevuQzS*#RI=C(RFQnUfSO0sIF^Xp}ox|yRPjp@YV)@ zS*mmGhR(}|WrutBmT-MUme7l4 z3B6dB&?B>iUOY?aC9;Gbl_m6&Swb(BCG^r+LNAjg^s-q(FPA0s@>xQ!kR|krSwgRr zCG^T!La&k~^r~4xua+hB>RCefWeL4Tme6YsJ{3pHq$|J({wX=j?CrjvcvxHtR zOX&5pgx(-a=nb=k-Y84xjkAQ_BunT`vxMF(OX$&ALT{cW^cGn{ZV#LT{fXw4c@X^Qw1D6MCmCp?A&_dY3Grcg+%dOqS5QWeL4|me6}-2|YGT z=smN9-YZM!y|aYgCrjvkvxMF+OX&Txggzik=mWEaJ}67*gR_J_BunT+vxGh@OX$P1 zggzom=p(a)J}OJ-qqBrQCQImJvxGh_OX%aXggzll=o7PqJ}FD+le2_AB}?d2vxGh^ zOX$pP41}_$;B%$`X1)me6Ns34KnM(C20eJuyq@Nm)Wq&Jy~(ETPZO z5_(FO&{MO7o|Yx_1zAF0m?iW@Swdf&CG;g(LSLFC^krE>U!EoO6GxvxNRfme99n34KSF(067DeOH#ycV`KGPnOX4W(j>?meBWS3H?Bp&<|z_{ZN+B z4`&JeNS4r#W(oaRmRK?Uc$Qc`{X~|~Pi6`IRF=?BX9@jGme9{;3H@A_(9dTH{X&+| zFJ=k-QkKv!XNeWkuVe|=U(FKwwJfn*`t>ZaZ2FBXq2J6B`mHRnOuC-qC(912cCOvf z>0w#oriJG`SL^xnazZbVCG>(>LNAmhbpG@9xmtVo10O=?`tU5_`iLx{7tIoSu`HoS zW(mD`me5ON2|X%H=q0m+UMfrIrL%-yCQImLvxHtQOX%gZgkB*_=oPbsUMWlHm9vCi zB}?d4vxHtPOX$_Jgzn1{dW|fh*BpE}@at;%LamiYgzIZ(3B691(CcOiyt_kQ zL6*=PW(mDfme3n#3B5^{(3@rny;+vfqqBtGJWJ>;vV`6;OX#h#gx)$!=xwru-Zo3< z?Xra4K1=8wvV`6-OX!`lgx)zz=v}ge-Ze|;FG=SwioTCG^-Vq4&%Zdao>@ z_s$Y}pDdyG%@TUQETQ+$68eBFp%2Uw`k*YK56%+$kSw7O%@X>sETIq268eZNp^wZG z`lu|SkIoYMm@J`>%@X>!ETNCj68eNJp-;>b`lKwOPtFqhlq{i7%@X>wETK=&68elR zp~qzjeP))>#t@B{aTh-BmH`o=u5wmCG?wFLcf(IR!?`It0#L=SikF9|G-NJe|g`+ zcJ{vqLG}41Y?#*@KKO9p*IV*?M|AbbuI}qk~CG>l;gnn<9(C^CMWtJ$rAbl zSwep>OXv?}3H{+Lp+Ax(^hdLV{#cgKzK6${FukM4uFY$oi;w5eP5ASFB1`B`W(obN zETKQ0CG=;qg#K)n(4Wf^`tw;re<4fgFJ=k-r7WSpoF(*EvV{I>me60z61qQ2=&xr9 z{f#W4znLZUx3YwunI-hMvxL4jOX%-p3H{wHp}&_U^!Kxb{y~<|Kg<&PM_EGuI7{g3 zvV{IgmeAK{34KGB&_B%*`e#`}|2#|RUt|gW%PgUPl_m7AvxNRlme9Y=68d*pLjOKX z=s#o$?YUsI34RU;*M@!;hJF_2wV#DC1K<5~-WUCEx9fdT*E52yX9VqhKJdPdOojG*foLDw^au4e>Y&j`An5p;bn(Dk`M*ZZQb_eEXri@M$yb-gd@ zdSBG_zNqVcQP=yTuJ=VUv+){rNf&y51Lcy)Wu|U)1%!sAuNw5W3zM z^|g69q3eB7*ZZQb_eEXri@M$yb-gd@dSBG_zNoLu$3bY%1MRs$*ZZQb_eEXri@M$y zb-gd@dSBG_zNqVcQP=yTuJ=VUv+)?u&gUOz-HiYeUZkv(LyrAMf<% zg8IGCpYH*}pMB_gVcvUQ7&GwIpYwUazt5*WFX(zk(DjU<>ls1M$lE1!JtOF=^KwGh zGlH&X1YOSvx}FhqJtOFPM$q+)pz9ey*E52yX9QiJ7j!)%=z2!b^^BnF89~=Gg05!- zUC#)*o)L6CBj|cY(DjU<>ls1UGlK5V*MZRWxj@%5g05!-UC#)5X5J2=>lr~`o0k*1 zo)L6CBj|cY(DjU<>ls1UGlH&X1YOSvx}FhqJtOFPM$q+)pz9ey*E52yX9Qi(2)dpT zbUh>JdPdOojG*foLDw^au4e>Y&j{LmVOJBTcl6k`q343xX9S;*clvX|(9iztbAYe+ zyzsL>%ny6$ANY0fm-f8ixek6$`+cs>mnQW5Swb(6CG>(>LNAmhbpG@9KbNH!NfUZ_ zme3=zgkCgD=*6;x9+@Td;#op3ktOt~ETNao5_+jDp_k4QdYLStm(3D-xh$cV&k}lt zETLD-5_+X9p;yildX+4pSIrW7wJf1m&l0*XOXxMSgkE#-;lQt3=L@w~9ucmuoh9@- zSwgRyCG>h(La(1C^afc%ZMemeBiV3B6yI(EDczeL$Ab2WAOe&LZ6={^pq^2r)CK~ElcPNvV^`cOX!QTguXaS=u5JMzBEhd%d&*NJWJ>+vV^`e zOX#bzgr1e{6XE_&dl`oQoxXGab6NiW*BUwT}nkDpOS>m4b<5|Lghx$a8&`)Lw{Zy9FPiG1JOqS5kW(oaV zme9{<3H?Hr&@W~Q{Zf|DFK3DWKm0p=`vyLe-18UlVD|r}S+(cQsSEwD*V=Q`wmn|2 zJ^%E2z4lzYU60r6^KQcSJ+6%(y!5UedR%MI=e_=&zU|&HET8-CKVN&jS9<*U+H>xX zJ>Fk?n>+Qm)}GsU?s2WX-Yz|^^>zMyyLMfx_k%G#u617RUAuK%Yxfk_ckjB^zILwd z(RHo;xm_FExi+5&eFMJ^{^H+rx@PG;JHK~S=VjBggnwt)yR(FTPnOW{%@X>3Swhdq z68imFLSLOF^fgℜ`Zf4`vDdp)8?4oF()}vV{I5}E=%a&X9@j>ETKIYj5fi~0pZ%v&%)5p!o2pgFlOMpf6n`&f3Ir2FY0wQtz`=YM*MP2WUy51Lcy)Wu|U)1%!sOx=E*ZZQb_eEXri@HBw2SV5TqOSKv zUGIy!-WT=Eyd6T<`=Y)!FDG=pFY0<<)b+lo>wQtz`=YM*MP2WUy51M{b@@05?RlU* z7wCFl)b+lo>wQtz`=YM*MP2WUy51Lcy)Wu|U)1%!sOx=E*ZZQb_eEXri`sp$&xGk6 zJ$7yAxnTAg+2`Y({#;PM7y9!(K=`u{Jul3A&kJJ)zWQ@MFZg%u*XIRY&j`An5p+Ex z=oxvtgsx`+^!HX9Qi( z2)dpTbUh>JdPdOojG*foLDw^au4e>Y&j`An5p+Ex=z2!b{rNf&x;_`^dPdOojG*fo zLC?(FA#^<>=xg(GLf12bu4e>Y&j`An5p+Ex=z2!b^^BnF89~=Gg05!-UC#)*o)L6C zBj|cY(DjU<>ls1UGlH&X1YOSvx}FhqJtOFPM$q+)pz9eyyD#i&!t{MM3&HtW(mDmme3=!gkC&L=q0j*9+f5Zl37A8l_m7jSwb(9 zCG@gcLNAvk^zvCkuaG75idjOhlqK}aSwgRpCG@IULa&x3^y*nc_hpG*|9^XH4qh7g zb+3G(*2*Ko>#v<9^g3BWubUdgm;mcgYgk z{lA?7^q4fEcgqrb_bj3J$P#*Nme6}<3B6aA(0gYIy-$|V`(_EfUzX7OX9;~kme2=g z34KtO&zA{VbtFnZimF*K@=>KnT&VM(Tzh8D|J}-p6D@*9RvxL4UOXz#E zguX9J==-yTejrQe2eX8JC`;&vvxI&mOXx?lgnleb4E_J@&H3-<`ak{u+xyAjr*q)k zJgoEB&(?f-Oz8QugkB&^=moQcUMNfG{AcLwJn-2V{CupC$Ar(%idjOhlqK}aSwgRp zCG@IULa&x3^y*ncuaPD6T3JG`oh9@-SwgRyCG>h(La(1C^afc%Z!vV=Y|OX#Drgg!b;=wq^kJ~m6} z;_LKX%U&dXFrj$7Tt=XO_@=WeL4^meBiT z3B7NY(EDWxy?>U_2V@C-V3yDaWeGiWcAj_7&c6PW2fq~tetosm?PdR-J}*}YzyJI! zp{HaCJvB?{X<0&FkR|kmSwdfwCG^EvLSK?4^rcxsUzR2GUA<=M{qweGmEN!G;}}=29nf{{+;VNzd>r-TZ_jaFdQsQ4_I&5sft|NGN`&bh zJzgl+-kl}%d$NRnZap|8mj`U6=)e=tkv4`m7c;Vhv)k|p#< zvxNRwme4+jV@#Od(PP)-J^XDe~`9F~*^e3}~{#2IGpUx8cGg(4^HcRNwWeNTH zETO-UCG;1wg#J>N&|l6H`YTyNe>F?!uVo3{pC$CyvxNRemeAkK68c+NLeI<+`rBDT zUz;WLcd~^3ZkEvB%M$wgSwjCHOXweF3H_rip?{nu^mSQ6|0GN3>$8NuAxr3=W(obX zETMm%CG;<{g#Klg(7(zO`qx=P|0YZ5-)0H@yDXu9pC$AkvV`_rFxmt^2ZU=wKMO-Y z3-j8~!kB^Y{yFc9!}GIH?~A&g5p+ExXy*e$&&U$Go)L6CBj{`L`v_gn2)dpTbUh>J zdPdOojG*foLDw^auFnO!J{RbEU)1%!sOx=E*ZZQb_eEXri@M$yb-gd@dSBG_zNqVc zQP=yTuJ=VwQtz z`=YM*MP2WUy51Lcy)SC_#Xb|Jcl6k`q343xXJns`clvWd{a)zL_WgW?>#S! z8TjhY`Mj`L*Ykp|X9Qi(2)dpT^o+b+Lf12bu4e>&O@1Gt>ls1UGlH&X1YOSvx}Fhq zJtOFPM$q+)pz9ey*XIRY&j`An5p+Ex=z2!b^^BnF89~=Gg05!-UC#)*o)L6CBj|cY z(DjU<`}27xbbT(+^^BnF89~=Gf}WYTL+E-&(DjUls1UGlH&X1YOSvx}FhqJtOFPM$q+)pz9ey*E52yX9Qi( z2)dpTwEM!YCQR?>v1|2=pog9dX8(OOpO1IibHUKh{_JyrukXC@vp>uad*~ndb?_H| zcFzl*>)`ja-{;zVX+qDRCG-MWLNAyl^g>xeFPtTG{`2<#Mph3`6M96J(2HgXy_g-S z!5*2%gzJlE3B5#?(4(@1UNTGQrLu%xI!owfvV>kXOX%gYgkC;N=oPYrUNKAPm9m6h zIZNnOvV>kWOX$_IgkC*M=)NqW*T@oj&B2EQzuufL6D!>tzYO zewNT1WC^`tme3nz3B7Ta(3@lly=j)vn`H?-I!oxyvxMFvOXw}Lgx)Gk=&iGa{-5mK zXSkL1y|4RVjo7ev9rmu+6&v>6dqqG*R6qfHffN$5(tA&43JK{ZCM!MJ2Yc_m#$Hp@ zQ?9+v*=K)0V~$6DJQtTQ_FDUmjQh&Xci!`va{_}o$MZ+T?3^X^E?GkFnkDpZSwfG< z5_CG>$=LLZbR^ubv| zACe{Xp;GzSwbI?CG?S5LLZeS^wC*DACo2Yu~|YNmnHP^Swf$XCG?3|LZ6f+ z^vPL5pOPi?saZmwmL>G*Swf$YCG?qDLXXN4`m8LWM`sB=CQInCSwf$kCG@x~q0h+@ zdVH49=Vl2#Axr4tz{jV&sSo+N@q2J09`t2;SXxiC%hzZjMr~B=Zt!v}++DX0ZN42hvEAMx7 z>)PmCYo1SzX!+LK!BTKB7YahrG`h!_Qe<(}n4`&HI zGfU`?WC{JzETLy*3H`Awq5myQ=#OU!{fR80Kba-;r?Q0pbe7P5Jx7`_qoa?_>%6-7KNMmnHP~vxJ_VCG-!ng#KZc&~vhc{!y0DKh6^RCs{)O zG)w58WeNTBETMmqCG@pfLjN*L=wD?CeO;E&zs?f+H(5gedzR3@%@X=|SwjClOXxpj z3H`?`q5qU6^q;eY{!5n7f6WrwbHQ*Ed>s(3^}iPSUkeM{*TTr|XaD^7#WnM_Q16Sn zo)L6CBWULXLeI<+x}FhqJtOE@c{@VaGlH&X1YOSvx}FhqJtOFPM$q+)pzCviuFnO! z-WPSfFY0<<)b+lo>wQtz`=YM*MP2WUy51Lcy)Wu|U)1%!sOx=E*ZZO#$m34vdSBG_ zzNqVcQP=yTo}G_F=z3q&^}eX*Uv+)^}eXwQtz`=YM*MP2WUy51Lcy)Wu|U)1%!sOx=E*ZZP&U+kDL zqopvIFJ0p8M-s{f=^>bk$p96%Sd;fW1;d@>f+5Od@|9N5VlJh?==z2!b^^BnF z89~p?$0c+JdPdOojG*foLDw^au4e>Y&j`An5%fSFcS6_a0$tAtx}Fhq zJtOGZ`8b5GX9WGjyq?f=vV^W@1YOSvx}FhqJtOFPM$q+)pz9ey*E52yX9Qi(2)dpT zbUh>JdPdOojG*foLDw^au4e>Y&j`An5p+Ex=z2!b^^BnF89}=*>}|q~o?dpXo)L8a zxnSNuNAq~R*PaXdU;FdU0e*cKhOhl0zSzS+_uIj5{M)ArrvxTl_3ti6^ zx}Gg`JzHpx;Ybr^^z^c83)^E+&sl!{^_->aIZM~)R$b3ox}LLiJ!k29&eHXqrRzCM z*K?Mx=PX^%S-PIHbUkP3dd|`V`E@7scd~@8=PX^%S^E2VJ3`OS61twV^bhlTLeI$( zx}LLiJ!k29&eHXqrRzCM*K?Mx=PX^%S-PIHbbYSX^_->aIZM}bmagY4UC&v%p0jj4 zXX$#*()FCB>p4r;bC#~>EM3o8+H=8h6Z&5Z{jUXI5AeCMFn%qJ?EdWMzb|f_$D-aB zbv+~KdPdOojG$-c;}W`_5p+Ex=vjF?Lf12bu4e>Y&j`An5p+Ex=z2!b^^BnFbAhh+ zMP2WUy51Lcy)Wu|U)1%!sOx=E*ZZQb_eEXri@M$yb-gd@dSBG_zNqVcQ4i#CCv?3p z>Uv+)^}eXwQss9%!Bm zbbT(+^}eXwQtz`=YM*MP2WUy51MH`(nq0 z89lx1TK~CV-k(u;Jl^Zi1@&`bAfE$-pL_p#VctKR@OUf?&kG~Fzxwk(FWj)?{Lc%z zo)L6CBj|cY&@=OK30=ls1UGlH&X1YOSv zx;`)HdPdOojG*foLDw^au4e>Y&j`An5p+Ex=z2!b^^BnF89~=Gg05!-J&?zp(Dk`M z*E52yX9Qi(2zqut4x#HALH{tXC-j^wq3anz*E52yX9Qi(2)dpTbUh>JdPdOojG*fo zLD%O3UC#)*o)L6CBj|cY(DjU<>ls1UGlH&X1YOSvx}FhqJtOFPM$q+)pxqbtHep6j zFS}OH2)h4VF!+4XJU;KW=Ysy%{=9R5&-aDlYk!DO<}lFxcJQ0#d7-&BDX%6@dcSYy z{rSa!>~VAN`J{Qji<=AClUwgsubtAmHc;O0*51eTYXKj-`S@e=ev70Dy=a!ui)9JD zc$UyhWC`uN>wcHL&a2+}yV6D!>tzYu$r5_~ zETK2d5_+R7p*PMFdXp@nH_Z}yvn-)E&k}lzETOl|5_+pFp|{QwdYdevx6KlIc$U!H zWeL4~me4z73B6;M&^u)by>ph(yJQKyYnITvWeGhZOX%IRgx(`d=smN9-YZLJXOHIm zp+}|(y-$|V`(_EfUzX7OX9;~kme2=g34KtO&zK1=9xvxJ_ICG>e&LZ6={ z^aWW$Ps|c}QkKw@vxL4dOX!QTgr1Tm^wcb&r)3F!ahA}RWC?v~me7}F34M8%&{t#$ zePx!=S7ix(b(YZAWC=YzOX#`TjtFCV?R6NPzCKHwmA)ZM=o_AzFIxEiPO?=W(oaP zme6l!iBr?{Y(7r*IJCJ9H@SE1e=OmJ`9J2=-nC8hb825}9{Yzg0Ll-7@- z{@PD%U7OJTID_B3*8PN-_j9iAH?8%4&ExAmFK%7)y5{;N3%J(t2H}9-wd)46`+LqI zy_fy}A1#uX3B72R&>qL;vD7{dx}OTY>wXElzl^=ko4oV;PH*qK@BR>a*L~hMpZl)+ zyl+1D_0nlVFOwy-uM0vimnHP_SwauX5_*L!p;ycjdZjF(SI!c8l`NrG%@TUGETLD= z5_*j+q1Vh3daW#>*Ul1poh+f(%@TUOETKDDqW}MIIREcMJwDs#aUt{$SwioaCG<{N zLhqa<^e$OK@0um_ZdpRR*EIJVy?dI_dt?c{XO_@=WeL4^me3=!gx)7h=zX(<-Y-k& z{j-EVAWP^2vxGhOX#Dsggz!q=wq{l zJ}yh>n<>v7eJATpyn$^to9=PskGbyey&5&l38AETJc62|X!G=*d|^ zUzjEIMOi{m$r5^MmeA9(guXaS=u5JMzBEhd%d&*NJWJ>+vV^`eOX#bzguXgU=xefs zo}MN2+-ygL{{J`l{J#(N|NnP?9v4DCkR|kkSwcUQCG^8tLO+rv^rKlqKb9r*<5@yK zktOt#SwcURCG^u-LO+uw`u~51^Z!2dUwM4q%;Q4nx3YwOJ4^gCUC;l`^NiPCmgl~D zZJO+HdF$Fhb0O>h&(Qn*s`t7fS>i)^zeTcyUNlSS#j=E6JWJ>$vV`{Cb-&AA=T+~V z$6a44P3WOnLNA>q^fFmOFPkOwa#=zzpC$CLETLD(5_-ifp;yWhdgUykSIH83)hwY` z%MyC^ETPxP5_-)nq1Vb1dhINs*U1uk-7KNk%M!YiCG`4PLT{KQ^hQ}iZ=5CcCRsvn znkDpRSwe4~CG-|qLT{NR^j29yZ=EIdHd#V%n4*U-akv| z1G0oZFiYrzvV=Z3OXx$ggg!J&=)=gg!P)=;N}4 zK0Zt66S9OpF-z!^vV=Z4OXyRwgg!M(=+m-U!EoO6(2r*c{X~|~Pi6`I zRF=?BX9@jGmYAOYTb8&w{cM)d&t(bye3sBJWC{IZme4O{3H@@G(63|({c4ubuVo4S zdX~^{WQl9if6o%G|07H2e`blR(*Md5SEk?068fzyq2JCDSETc(F7WS|xBox4$9dBp z8xHZpy!WfuKHU14qw|){^U2KCweh`=XFk%pHZ8B4+PnVI)-|to9kW{3HW++$_uG%P zu6bprdnNzAcwG7PiOYGx>QFe*T}$68dvlLVrF>=r3dm{lzSyzmz5Pm$QWaN|w-H z%@X=+Sweq3OXzQ83H{A1p}&@~Y&j@-}-j2}qjG*foLDw^au4e>Y z&j`An5p+Ex==xls>vMsw_eEXri@M$yb-gd@dSBG_zNqVcQP=yTuJ=VUv+)^}eVF^0*Va-WPSfFY0<<)b+loXXoP(y51Lcy)WuHc{@Va`=YM*MP2WUy51Lc zy)Wu|U)1%!sISfUOK8sn&2xdS_eEXri@M$yb-gd@dSBG_zNqVcQP=yTuJ=VUv+)^}eXx7ds}*=;>wG`p*UP&d45*_xf`|{ahHx=K$g7-hWBcM$q+)pz9ey*E52yX9Qi(2)dpTbUh>JdPdOojG*foLDw^au4e>2 zkjI_S^|?UTGlH&X1YOSvdUiezq3anz*E52ileZ&uJtOFPM$q+)pz9ey*E52yX9Qi( z2)dpTbUh>JdPdOojG*foLDw^au4e>Y&j`An5p+Ex=z2!b^^BnF89~=Gg05!-UC#*G zePM4CX7u#3YxRtv`_BdQ{yCb*cMkCDyD)t15Anqw2D;x4e&grvdBJlX zyiN0Vt}T)#^rBfpFP0_r;#op3ktOt!SwiQ3U7KDiFA%N|%@TU)ETNal5_;Jzp_j`N zdigA&hh+)9LYB}gW(mDgme4C_3B5{|(5q$%y;_#gt7i$lMwZZPW(mDkme6Zw3B691 z(CcOiyWky%0?l_m7iSwbI^CG@dbLLZkU^zm6jpO7W=iCIFQ zlqK}ZSwf$ZCG@FTLZ6l;^yyhbpOGc>nOQ=Q$`bmlETKnd2|XrD=&`n2gMD^hCR`tv zCGLf@Jt^le!}-<~D(9a%!(nI-gHSwi2PCGSwcUYCG;a%LO+@%^kZ2AKyxAM@7;S$njgRF_$K^U{2l%J<8S|+{N_2`_jT~%zMc1Lqndl(S6e@R*P36u zuPxwO^L}4nz_pHd4u4<7xAVTgX|l&RTJKk{ee+$eeXDh?xo7&Y-)>#=y5{;o@7i@< z8GL>c{{EUDYWIiR>qWDK>x*Rx?V;>GoW1sI+x+^vp8vh``#tTu@BR>aZ|`%s`8;;r z=Wz3Rtd~v`dYLSteU3Mu=X$v`p_k7RdRUgwD`W}1VwTV=WeL4%StlO^=FSwaub5_-EVp|{Tx zdWS5bcgzxcr!1j&&Juc;ETMPJ5_-2Rp+{s1y?d6>dt?c{XO_@=WeL4^me3=!gx)7h z=zX(<-Y-k&{j-EVAWP^2vxGhOX#Ds zggz!q=wq{lJ}yh>T`ob)sFUk^nN|w-5vxJ_OCG^EvLSK?4^rcxsUzR2G&vV?vzOX#Pvgnl|p=x4G-|L>>I|NFiF?B7pso)zhbdRCUuAIlQ@-?D`Mc$Uzg$P)UK zSwep*OXyE$3GHzjX~K-2UUqF^dn`VapEu#>|Jf{|KbIx+=d*#FTUkPXJ4@(+ETO-XCG>Z*g#KQZ(BIDzdUlr3 zKgbgLhgm|;$rAcUSwjCfOX#0u3H{S7p?{Vo^v|<|{zaD1*JcU*%PgUPl_m6bSwjCh zOX%NZ3H|R`LjN{P=-*`t{rfDT|Bxm0AG3u1Qls1U=K@`y3v|6N>Uv+)^}eXwQtz`=YM*MLm$mozV5ZsOx=E*ZZQb_eDKBABWKOzNqVcQP0WS5xU+Nb-gd@ zdSBG_zNqVcQP=yTuJ=WKZN6VZdmd<>3v|6N>Uv+)^}eXwQtz`=YM*MV+7Z^o|KLdUjvVwf=L#yfd=L>b?G4P(K$2@<~AWx%ZzJ z7QW|&k=ls1UGlH&X1YOSvx;`)HdPdOojG*foLDw^au4e>Y&j`An5p+Ex=z2!b^^BnF z89~=Gg05!-J&?zp(Dk`M*E52yX9Qi(2zqut4x#HALDw^ao|CsDbUh>JdPdOojG*fo zLDw^au4e>Y&j`An5p+Ex=z2!b^^BnF89~=Gg05!-UC#)*o)L6CBj|cY(DjU<>ls1U zGlH&X1YOSv+I?Yf6K3@EvTOB>p!?4S^Zq%S$K$>BT+sj8pLY)M>$@;~?GN$A9tOJK z4u0e3?s>s;9lTBRcCIawCiJ3NLNAsj^x|1UFOenml37CMe_fkiDlZVO56u#K=`5j_ z$r5_mETNan5_MZme8wa3B6jD(5q(&y+)SMYi0?( zR+i9fX9>MdmeA{F3B6vH(48!y*Uu7qgTWWO-_A~Nn3o9GH_8%v<1C>!$r5_gETK2c z5_nx$S$r5_oETM;I3B6sG(A#GTy+fAJJ7x*JQX9>Ma zme9Lq3B6mE&?B;h-aSj`J+g$}GfU{bvV?a3Z_WUEWSY?XWC^`*meBiU3B7-o&qq2lPD@*9nSwfG=5_)Ww z&}U}}JuXY=bFzdUpC$CUSwc_968gL>q0i3}`hqN>CuRvfDNE?dSwdf!CGzA{VbtFnZ?I!owlvV@+VCG^~EM}!ad z+UxLv^z~W7KR>!5OXwT3guW?D=$o^Iz9mcOTeF0|ElcRzvxL4QOXxeZguW|F=)1Fo zz9&oQd$WYTFH7kAvxI&iOXvr)gnlSX=!dg}ek4oiN3(=}EKBIevxI&kOXw%FgnlYZ z=%=%UekM!w|NFP||NA{x=J9zmj|-vS$`bnREYbe&-;N)Ab@$sJUAt~D-!~*n`0v;*k|p$_Swb(CCG_H1 zLNAde^paUZ=YL(BUMepTt`E%;dg&~om&p=(*({-#%MyC|ETM;G3B5v=&?{yMy;7FY zD`yG4N|w;8W(mDome8wb3B5*^&}(K1y;hdcYi9|)PL|N?W(mDsme8Fnq1Vq6dV|3i zyWjpKy?^LSK|6^pq^2r)CK~ElcQ&vxL4ROXy3pguX0G=*zQ&z9LKL zE3<^YDog0AvxL4TOX%rYLeI^1MEG^Dy$;ur z#Vny;$`bnJETLb?68hCFpZK-^dczrvIKLT>nRw(ErR5KTrQFOZ+VTW|q)z zWeNRumiTGfS$z=`&W`7HU&^29<(i({>)-U=c1W)`=&ku}>zem#_+9H-eZSwgu1#p( zDf@@kwfcU4>|I;3ynUzl@}F~Usq%XNEuNvRUbfYpRyWT-UVU@Eley*kpL*YScY3eC zu(#_hGXMAO@3zkWOX$_IgkC*M=ryv0UNcMRwX%d>J4@(wvV>kYOX&5ogzjVsy?&O^ z8)gZ;QI^mfX9>MYme8AK3B6gC(3@uoy+xMLTV@HpRhH0OX9>McmeAW~2|YYZ=$z_stS|zbv8m&l388ETIp~68fMlp%2ay`j9N456u$#uq>eu&l38GETNCg68fkt zp^wfI`j{-CkIfSLxGbTM&l38CETK=#68fYpp-;{d`jjl8Pt6kgv@D@d&l38KETPZL z5_(jY&}U@{JvvM1FN3&`c~Vm!M-go z6RzK$CG;IxLf@Gs^j%p(-<>7&Jy}BEnBUwT} znkDpOSwcUaCG-=vxI&nOXyd#gnlhc=-0D^ej`iVoc?>3aQz=yLjN;M+?f7XmbfAPW|q)zWeNRu zmbgA$&*o!gk2_n}rZpF`ceSp$G^Ka_^~Jp~@6T4$INyZ3-{t+jyrj37_kNRQk9%70 zSFhdsF4yjBU7OT=!0i34Yc4g{A6US(jyDJo_OAVUFuT9E7}9&$f0JR6yiDjtvxHtO zOK1;Q_aW=`5?R9a{O_ILUufTb_lMBCzEqZQeQ1`@OJ@naOqS5gW(mDqme9**2|X-J z=oPYrUNKAPm9m6hIZNnOvV>kWOX$_IgkC*M=ryv0UNcMRwX%d>J4@(wvV>kYOX&5o zgzjXC|46T&C0yTN@Wt-82j&-Q!@NYezEPIY8)pfOX%IQgdUM4 z^zK@Zgl_m7vSwfG@5_+F3q4&)adcQ28_s@14Fj-_sZ73gPW9 z$`X1?me5nP#6$W1rez7&FU}JBk}RPw%@X>uETJ#Y68efPp|8vm`l>9Uug((snk=EG zX9+zw+YzDv-(Z^m_aFY7Mfc}%A@l=TLO+-#^g~%fKb$4>BUwT}nkDpOSwcUaCG-srSvdpz5_Hmuh(vV=dQ`#_e^AIuW^Ls>$9I7{f6Swep# zOX!bg2|X)I=#OOy{cl-9e>_X*Ph<)G$tb& z&;PSoLVqqx=+9>f{e>){znCTTm$HQZa+c6v$rAdjSwep;OX#m>3H^;Mp}(0W^tZBv z{&tqo16e|UCrjw>W(obhETO-jCG_kpp?{Di^bfOyo|7f?kFtdRahA|O$rAdfSwjCT zOX#0x3H=M(t--!FFB7i+GE3-RWeI&68bk;LjQY~(7(+R`gd7E|2|9TKV%90 z$1I`$lqK|^vxNRjme7CA654aYa1(qT5U%yV7W!Wc3)|Pi$nIzV{P#uw9#g$9>Uu`d z^^BmM4+uRoOXzw=(DjUls1UGlH(q1-d>L=z3q& z^}eXwQtz`=YM*MP2WUdLWNGq3eB7*ZZQb z_eEXri+XlG4x#IPQP=yTo|CsDbiFU?dSBG_zNqVcQP=yTuJ=VwQtz`=YM*MP2WU+I_KO!i=6? zcCG(hFz<})@p!L47u3&%fqV`Se(wF}g@x~VVPyALfBxr%kF}l`bUh>JdPdOojG$-c z;}W`_5p+Ex=vjF?Lf12bu4e>Y&j`An5p+Ex=z2!b^^BnF89~=Gg09aCx}FhqJtOFP zM$q+)pz9ey*E52yX9Qi(2)dpTbUh>JdPdOojG*foK@a3{Cv<%-(DjU<>ls1UGlHI- zk3;BsM$q+)py%Z62wl$zx}FhqJtOFPM$q+)pz9ey*E52yX9Qi(2)dpTbUh>JdPdOo zjG*foLDw^au4e>Y&j`An5p+Ex=z2!b^^BnF89~=Gf_7in+k_cCz3f^&Bk2Bf!MuNt z=J9y1Js0%9_UD}g{Q52oU;9IRv4?@~w}apKxqDt{u8m7wGI;5I&n@8E=-##GTh}_h z*P9nw*L=)|7hBgl&4uhst!vZDwU>L>mh8Q5sdCNV*BRRCWn0~8^@P%s78(3=^NGFJ zy_%0Vw_Jav_2U{S*Zln-k6%8-J6za%$=}zx=$+T*-z9vPkK^y=w0|9%ALHw-Yt6^^ z9&fa+HNQTt{k?VVJg@ArdiMdGF@0FCXJiTguHpx>g#KWb(DiJn>)BAx%-a#Vo(*+9 z8|qnkJ3@afOXzww)b(tr>)BA(v!SkMLtW2?+G9V`gc&`(?ApTiSbQcwZ^F;No|APw zC+m7|)b*UK>p5B1bF!}IWL?k6x}KADJtymWPS*9Dtm`>h*K@M2=VU#QpEseulO^p5B1bF#iR-w&bd zIa$|xrLN~>UC+t7o|APwC+m7n*7cmM>p5B1bF!}IWL?k6x}KADJtymWPS&0ahMVB) zfN-tUv+)v-5EXUGIy!-WT

iS%u>wQtz`=YM*MP2WUy51Lc zy)Wu|U)1%!sOx=E*ZZQb_eEXri@M$ywfkbngc&`(>{|c1VDR~%c|6|h&js~!VIV(g z!skN&d12l^Tk?1;49^QAyTAJLKQH(@zV&%Q*E52yX9Qi(2zq8dE}`ogLDw^ao|U&F zbUh>JdPdOojG*foLDw^au4e>Y&j`An5p+Ex==!{%>ls1UGlH&X1YOSvx}FhqJtOFP zM$q+)pz9ey*E52yX9Qi(2)dpT^gteWLf7X4UC#)*o)L6CBk0-rIE1ce1YOSvdQRSs z(DjU<>ls1UGlH&X1YOSvx}FhqJtOFPM$q+)pzCviu4e>Y&j`An5p+Ex=z2!b^^BnF z89~=Gg05!-UC#)*o)L6CBj|cY(C!O+n=qrNmtE^W7tA{&cs$|vn$?cg`f^MdC(c$?;vxHtHOXyXzgkCL6=+(1? zUL#BBHM4|XD@*9LvxHtJOXzj8gkCR8=uVc<>t_kQ!QhMCZ~u|rFfS3VZ%StlO^=FSwaub5_-EVp|{TxdWS5bcgzxc zr!1j&&Juc;ETMPJ5_-2Rp+{s1y?d6>dt?c{XO_@=WeM&6-<$#T$TXq%$r5_sETQ+y z5_zK1=9xvxJ_ICG>e&LZ6={^aWW$Ps|c}QkKw@vxL4d zOX!QTgr1Tm^wcb&r)3F!ahA}RWC?v~me7}F34M8%&{t#$ePx!=S7ix(b(YZAWC=Yz zOX#`TjtKq#&e!~Z&)z>Dx<8K#p&!T+`oS!rAIcK?;Vhva$rAd}ETJFE68iBhp`XYS z`pGPzpUM*Y=`5k2$rAnl&e!~Z&;EuyK5yo6A@o~WLcg6Qu21_t^$-)9uUjwwvvqBR zSN6DR5iiVpKh^OK7W2)#Yt3_y*Z!;ZG3&KA-{sm{t!vHmyN~&H>zda!&tL!N0zjuD$pndnJxLNA{s^sp?USI822#Vny$$`X3zETLD)5_;7vp;yZi zdi5-!*T@oj%`BnU$`X3*ETPxQ5_;V%q1Ve2x|1dP|GlpHf8XHo**=d8p?Am(JY}K%M$wW zETNys68gz3p`XeU`spm8pUD#a|326Jzi;?o{qJ*KH~4UN&&@-6FZ&h&lugVho>MWtJ$r5^cme6yv9TECx=l$&n)p~ z`d?Y1e|GNwbLalqdGKq&yMv#{9pn_q(Kg%>S=-ZB)7TKU&wCdx1Ck|6A8Q zx4uL70h;$?@%m2h;*i$0NqK$qaTi&@wF`Rhw`l8HbDwf~vDP&ovtjYpHLvgVE-ul! zHnsP9vt;YqK)KfZ|ExE^Mt=Mgde`6IdcTXy`z_VFHa^#yAM4Q8wbAAMmTp~hA8lBs zb*+;N#5Wi7?Yv)y`fIyf>;0N@srOj^U9Js#muoAuu1)j7daT&G=GD#hl@@SqLht=n z?p?dC*B|K~v-oZ#{3|9Q~7zk}iN z*gTI5p|{8qddn=Kx5^TF>nx$S$r9S{X9+z#OX%&ggx)?&=pC|z-Z4w){_kn~e;#zg z;HN|P^Wnt2O!)jbDNE>+vxGh+OXyRxggz}x=+m==J|j!$GqZ$tk818y`m8je-4_Wx zCQInCSz?vkJDd{;*WIfLeNL9ph(LU*!6|KFwRpPl!~=K%2g2hsH`@`Re|GMlouAF;0pV==T$a$! zX9@j6me4O|3H?%*&@X2R{YsY5uVxATT9(kSX9@jAmgt|I`)B9=*?I75!E^f-d6Cdt zW(mDjme5;g3B660(A#DSJv>Y3?Xra4K1=8wvV`6-OX!`lgx)zz=v}ge-Ze|;-Lixp zktOu*SwioTCG?(ILhqF&^xj!QkIWL*me7Z134KJC&_`wo-9I}oyk}>>{*&{0Ki=`+oAWpj{#~S7vV^`hOX%COguXpX=sU86zB5bcyRwA7J4@($vV^`jOX&Nu zguXvZ=m)ZdelSbuhq8ozI7{e9vV?v#OX$b4gnm3r=qIv-elknwr?P~8I!owhvPA#y z@c!5QJ%k~H59d3-N0-h2e^+yDp!YV$ze`pBeU;|^?zL;=Iip@1Cwr{fy4E~zxx7~E zTK$-7x2}zDK45m8)-{)!>+80zdERWE3)fq~wM)En=(MiYpTqjCYt3`9_uHU#?QE~? zv0>|4{rPRwx;8%VGpYA}8@H}CKR1^*S-`dCxqH*rwdUvT{Wfb|t3Q^_Th}_?p~sP} zYZIFb*)3YvI_27y?{aOcce%E8>so!Rw`pDT-8H{X+b-Z*^LP&LUAwOL68gm~pHk?v|Li>Y zwcxpZi@ZqaEwhB)Dog0CvxMFzOXzL0gdUzH^mbW7Z=WUf4p~C)m?iX1SwioeCG;*? zLhqU-^ln)~kH`{w_bj3J$P#+bETQ+x5_<0}p+{y3-9J0`|2?JAgP#uF&xbL2neh4X zUOGGb`TXxZJ8z$}dH?L(|Ff2d2fuu}f7bCxUMBq6$D>(7Kb9r*<5@yKktOt#SwcUR zCG^u-LO+uw`hV8)U-M@z{eSK}cyCxCpHqbA))ljaUMWlHm9vCiB}?d4vxHtPOX$_J zgkB>{=ryy1UMoxJwX=j?CrjvcvxHtROXyCP=>K!){@J;IcJ7~@`)B9D=R?nrQ}QCQ zL(alevxMu@vV^`kOX&XDxqo)ze2O=60= zy#GE^z4jmPa_xVzcNgBC*N46CgN6ekxZ4oi3GNK;?(QJpbyp)GCNq)FRI zCHE3{+%3egc>aL_6Z8oV~w! zZwJhG&G*gz{0{sL*JhuygJ!OM+8czAMZ$e>!IsbqwS@MMJN1t|(~Go(>pqWWKbLfW zP5%6nCnr1~lix>9{yR%MOn*2`eLU>gW5UPBPA#E#ZVA0hOXyu&LhsfRdiR#lb6Z01 z(Gq&kme6~(gxLpU@Ke#Fo$}wS+#oCG;sRp-*iIeOgQC z(_2EH(GvR1me6Olgg(0^^f@h|&us~PUQ6inTSCuk34K9J=nGpyU(^!X8DsXGp`8I{ z&jAPjkJtYBCD;FV{yP}&fAagt$?qd4zmN3UIsNC2o^J`iuY92;^ouQ_Uup^ca!cr+ zw8Z4kFTMKvRj)mZ_nZ)ViI&iBX$ie#OX#<@gkGv8^wKS%C;txSO>bIz`stW@KHfI- z*zCL5o8GieOL)H4Z3(?zOX&4mLT}I#dc&5`8?}VqxFz%^EulAU3B6fM=*?R~Z_yHZ z%a+h@Zwb9sOX#gzLT}R&dfS%J+qHz=z9qDC^W^M2^>j=<9|!c1@cbOu68fN)&Zgg&q(^g%754{iy4NK5F+*?IDH=S!wP9Hu@V{yWajGw*Kd z{f%d5A8%K8kV^mQ$vuWt!`LrdrzTSDK|68h$r(6_XNzO^OvJ?$J3lmFf9 zGxYg|DYxG4_iY2s3r7|TSC7c&&~@>|McE;=ggYf|6bCkX1>plnf-Cc z|Bu)Hc@Di$*9pCFOXx*fLND49x~KM)KhL2Ts|np-ulyM*y?9ONC0atir6u%|Eur7q z5_+kY&`Y<3UZy4VvMr&PYYDx4OXw9^La*2odZm`oE4PGprk*`l>+xKz$8)tF&((T7 zSL^Xyt;chMUyPv|k>`iU)} zPihH$a!cq_T0)=N68f~3(5JVAKBFb{nJuBuY6*RIOXzc2LZ90b`n;CV=eLBO*An`I zme3crgubXH^u;ZqFKG#VX-nwKT0&pm68egk&{wvEzN#hk)h(f~X$gI8OX%xbLSNq! z`i7R!H@1YnsU`HyEun8|34Loz=zH2ZB95DB--l!CFSf)n^_N;gf4L>}S6V`UwI%e| zT0(!lCG-#UA zxpwMrxOU~twbSOiHv2X%pFZC;?{k(j=DT*u?1lEs`L2yWzq97M=I1y2&vN!}xHkK4 z&zbMqye`a{xp?k;*LIqHqxQV{u8sdJ=g(aG)Xa}xpd~Kq+JY^i7itN;a7*Y#T0$?{ z68g<8p%-fj-CwW#UPmup6MBi3&~Ir8y<|)1x3+{{swMQ&Euoib3B7Df=;c~MFW(Y+ zg_h7OwuD}(CG^TIp;u`My=qJ7)mlQY-V%C^me6aqgkGy9^qiK^Yqx}6rzP~dEuq(I z3B7(x=nYyzZ`cxgqn6Mcw}jrLCG@5(p*L#@y?IOMEm}fv*%JEgEupt+3B7eo=xtg; zZ`%@jyOz+~w}jrICG?Ihp?%)Ye(vg>YeMhR5_;E`(7UyS-n}LC+?LRLw1nQXCG=h` zq4#bHy-!Q%eOp5B*AjaFme2>Zgg&q(^g%754{iy4NK5ELTS6b!68i9#&_}d{KC&hB zQ7xg5ZV7!%OXy=;LLb)>`uLX6C$xkq0ep!eNId0b6Y~6*An{tmeBKBLSN7l`ofmb7ul(r?u&a&xPD1X=u2BdU)B=( z@|Mt7w1mF0CG=G-p|5TUeN9W~YgFOH1fmTSDK{ z&Jpp=nf877M*YQ>_0(2uu-exfDxlP#g2 zY6<;xOXz1>LOeWqvM{v9)Y<;;)2ZKki9>37x(^q9E5Ua%$fLM@>eZVA0eOXx*g zLch5s^kOZc`|Fi+xn8^`^b#$h-_jC#$(GP>Z3(?pOX#IrLNC)2dfArH%e91Fz9sYu zEumLz3B6KF=#^VSuhJ5F)t1n!wS->1CG;9Cq1S8)y;e);IW3{rZVA0kOXzi5La)~n zdi|Er8?=PpuqE_HEulAV3B5^6=uKNfZ`Kle^On$Cw1nQWCG^`{LT}X)dh3?Z+q8t< zwk7m-Eupt>3B5y0=p9=^@6-}{=a$gBw1nQZCG>7Bp?7ZyJ+~$F9xb8wYze(rOX$5@ zLhsWOdf%4N`?ZALza{hmEujx=34KsY=!08AAJP)~(3a4LwS+#rCG-(3p^t0{eN;>6 zqgz5B(-QjFme9wwgg(9{^a(AYPizT&QcLKQTSA}G68hAZ(5JP8KD{OM87-mDYzci< zOX#y(LZ8zTx7O#j#4R&@`BdHZh2N^r>oMW&&uz3xMAo^mwnex@2;=z`w4G8{?FSUgJa!cs1 zw1oa@OX#n)g#LO<=x?-y{$@+)Z?%N}c1!5*w1oa{OXz!BLf_XC`u>*C5441SuqE_E zEukN73H?Y*=to;ZKh_fZ@s`j}w1j@LCG=A*p`UIE{Y*>fXItXx`g<*LW&K=B=*rq$RGZf7%kR|EwkS&s*Y(`WG#6 zdHu_l(7$R4{p*&vtRBzpJINV-d%kPFd(Ov_H zWm-Zn+Y)-Yme9+$gkGT~^olK^S855pa!cq{T0*bd5_+|k(5ttEUZW-Snk}K%Y6(53 zB_@CN{TEuoKV34MG^=o4B(pV$)mq?XVp zw}d{WCG@E+p-*cGeR@miGg?BQ*%JD!me6On#MZrVpVJbqpW71pyq3`Cw}hV868eIc z&=JskeQQhTd)hf7CV%$+wdZmFyVD1HUI_hQOX!DMLOJcw#4MWyZ_qr_%C{Xe%bRv=wG#j{&h?I zymlV{w21{~^w{Utf-RvJY6-n?OXx*fLND49y1yo8=c#97`uSM7$Asr+m6p(}wuD}- zCG_eoq1R{$y=F`3wOT^YX^F|%d2)8%qmKi^$H$&6q4#PDy?0CKeOf~A+Y)-eme7;4 z^XupA?AP<0eq9M?>2q7c_48UnPtMMhv-8)dKOCmccHihR;jH(~meAj73H|Ms(BEkZ z{oR((_qK$-uO;;TEuo#82|YPGzxuW8$^UNEN5yac@1F15YoBXNo49R8kNvs-ceI3l zXG`dJwS<0mOX%BMLjP7v=-+M${W~q8?`R4Ao|e$R+Yzrof7KHDuUkUj)e`z|T0;MAOX$CA34M1< z=)Z3X{SPgnKi(4hA6r8IQ%mSiw1oa-OXz=Y3H_;-(4TGz{Vy${|FtFbzqN$^OiSp` zwuJs%OX$zHg#P!I(ErgA`afGj|5r=sFSLaATCk-Fp9h3%lh1|8=fWH9b7Ah(NB?W@ zi~hHf<9$(&X9PW-5w!CGp>J;qJ)RNtct+6U8A0FC+Yx#^Bk1vrpvN|;@clCYD1fcwf}xeNo@t+Yx%aFY58WsK@)F9`B2K zyf5nUzNkOh--pm%2WDRj^mt#?<9$(&_eDM47xj2w)Z=|okM~7A-WTD1fcwf}xeNnqF&M|S@j2^o-c`bPLjO_V%y}uTW9}9Q(F+ljaPhJ<^c&`g{r+)Qc z`?}!2modIB=<$r8$1{Q+&j|YV-Y=oYGlCw^2zopt=sS8lLXT$zJ)RNtct+6U89|R{ z1U;S+^msx7i|gs=9bWlwS?}kSN{CFUc4ss5-p+M(h_>f zme6l)3B6QH=%rgiFVhlw*_P1DwS->2CG-j{p;v4Py;4i)m0Lou(h_>rme8xUgkHTR z^cpRp*K7&BR!itPEuq(L3H`R|hg09ar(UN=gzM|JgkG;D^!hEKH)sjHVN2+ZT0(E! z5_*%C(3`e|-mE3`<}IPOXbHV#OX#<^gx;zp^wuq*w`mEzZA<9wT0(E%5_*T0&^xw- z-l-+@&Ml#LX$ie+OX%HNLc9OZo&ogSn$UZ+gx<3y^jiLO<3L`tg>~Pqc)7vL*CWEuo)o3H?k<=x1AE^54vS?Vo@D z|N7s|eB1QXI`!@E_FVh@_5w5CH~ShPfA4>C?f2)qHgD$b%^%En&G*gmP5j|}*T&bq zKl%;Vrr)aH^Uil|_O;XXKbh~^?CYp&e>&f_+1FIp{%q#jr$l%TW}kgkG{G^jlj(FVzxy>6XyTw1i%^ zCG>JFp_gw7y+TXq6&09im(Gq&gme6l+ z3B6TI=&f5qZ_^Tb+m_JVwS?ZjCG-w0p?7Qvy;Do*om)ch(h_>tme9Mkgx0))M;g zme5DEgg&w*^ieILk8TNlOiSovTS6b#68iX-&?mHnKCvbANiCsIZV7!#OXyQuLZ8+W z`t+92XSBqf^_eZ<`dKZZ&u$5QPD|)>TSA}L68ij>(DPbCU(gcz!j{k%wS>O7CH}m= zq$OOxv?cUqEuk-O34KLN=qp=7U)2)&>Xy*gw1mF4CG>SIp|5WVeM3v=8(Tu()DrsU zme9Afgub;U^gZny5g(mt--nOXUu=mF*I#N0{pFU>Uug;b)t1m-YYF}JmeAj53H{BM z(BEna{q2^}-)RZ`-ImbzwuHW~CG`C*p&w`o{a{Pzhgw2E+!Fecme7y3gnq0g^y4j| zpJ)mFWJ~C$T0%eF68f2z(9gEShwAUO#0To%!TP5y;rh>7LjSxa-e3QsCEi#6vL*DdT0;N2CEi<)zfYJ~ z&GB9F7r)`!^p*Zy?R?i}UmslmtC?$y&ivr1oOyg+vG~kmqnDho=gik{8~Uz({Osqp z@BZtlcYfQgpPKnG3$(;1y0&0T=!IHBFWeG(k(SVY4zoWG{pOm`i?xLAuUF3c+Rt_N z=c|{f3H_Fq(4LFg=R?1>CiGG*p_gt6y-Z8!Wm`fo*Aja9me4D-gkG^F^hzzES8fTt zN=xWfTSBka5_E#H*E>MSxe~6TS9Nq5_-#)&~I-Ey;V!-ty@BG(-L~ymeAX^gxZgg&q(^g%754{iy4NK5ELTS6b!68i9#&_}d{KC&hBQ7xg5ZV7!%OXy=;LLb)> z`uLX6C$xkq0ep!eNId0b6Y~6 z*An{tmeBKBLSN7l`ofmb7qx`GxFz%@Euk-M34K{h=*wF|U(pi!%9hYqwS>O9CG<5d zp|5QTeO*iF>svzK&=UH_me4n~gub~Y^eru+Z*2*EPdi7%Kh3o7!#~zvY>AK8Uup^c z<(AN2X$k$+me5~o3H|k!(BEhY{mqup-)af{?UvBrX$k$^meBXMgubsO^!+WNA7}~v zU`yzST0%eE68e#r(2ur+eyk<*<1L||XbJsfOX#OsLO|F9+Wk6J?ixFz)SEumj%3H@SA=$BeTzuXe~CoS>!^-o*E^`Etb{&`FM zUHyxe_}luIEunwa68hII@i(=z{#+Bc&FJyhyY`Nj(C=&s{jQeK?`{cwdrRowY6<<@ zEunv>CG;IFq2JRI`gdDG|6WVz-){;12Q8uhuqE^#wS@lTme7CF68cYDLjPGy==ZjS zeqT%I_qT-pKuhQkwuJsrOXv@`g#JiN=#RF9{#Z-sKW_WTox-k1Xw?H2Q#6^89EZ7oyp_b4Kw}f7#CG?^# zq2JsRda;(!{q@S%)AZssp_gb0{g#%{OSXi5YfI>*T0$?~5_*}I(95=jUalqd@-3lP zXbHVyOX!taLa*EsdX<*YtG0w*ttIs8Euq(F3B6`Z=(Sox&uIz0c1!5DO+TFewqJ;K zdPKOsZcFI(T0*bi5_*G{&>Oad-l!$?#x0>YX$ie)OX$s7LT}y@dW)9OTegILdrRo8 zT0(E#5_+4K(A&0z-mWF|_AQ}zXbHVzOX!_iLhsxXdY6{ayS9YhttIsCEurVOgx;ei z^qwuD_i72fcT4DfT0-yJ5_-Ru(EGQ9KA=p4Srkf|k%1wuHXOPStc@++)J^OIku- z+7kM*me7~CgubFB^p!24uWAW>bxY`LT0&pj68gH9(AT$wzM&=bjV+;XY6*RFOXyo# zLf_gF`kr=aZ<}Ca`&kLa+Yzh5POX!DNLO;?H`q7rqkF|t;ye0G#Euo)m z3H?+{=%-skKhqNW*_N36|2n?*>uX=_`T1qf3!#7268hII@s--|qZTl+)AUSz`^ovP z&AtY?_`XGa`|8`xz8?7AH5Z+}wzP@cX7t$qzs);ZLcg;m^t)O@zq=*$?Jc2yt0nYr zw}k$kme6;!gnmy;=-+J#{d+B;f4?R4AGC!2!o6%#}-e}LoANBJl{QUp8 zCG?-Pg#Oc((0|qv`n@fo-`5iQ{VkzC&=UHCEulZu68gg}p+C|R`lBtOKh_fZ&s##@ z*%JCMT0;M2OX$C93H{eCq3>!5{WmS4|F$Lc-?fCkyCwABw}k$Ome3z>3H^^Pq5r8R z^e0+Ef3hX?KevSbR7>bjw}k$emeBv&68hg-LVuhZp)$NQol?~8i8FY58WsK@)FzO(0<(BpkkkM~7A-WThZp)KiS`h&|U{-Ukmhj zU)1A$QIGdUJ>D1fcwf}xeNm71MLpgZ^>|;@<9$(&_eDM47xj2w)Z=|oyD!c$aoda@ zyEb_(c=e3z`FOp*7K|SYclI$r__37!tp0j^nf3F2=LNC-3df}GPi?oDZv?cVLTS70^61u-$ z`LjfN@tV*}w1j?3OXwwALcg^o^inOMmu?BYOiSoxTS70_5_OXc-nb?7CM}^iZ3(?uOX$s8LT}L$ddrs3Z*K{`RZHltTS9Ns5_;Q~(A%|y-o7RD z4lSW~Yze(nOX!_jLhsTNde@fFyS0RNewaN&=(#na_h<>dXG`e4T0-yL5_+GO(EGN8 z-mfL}{w<*oXbF8_OX!1ILLb}``jD2;hqi=1tR?i}EuoKS34LTs=%ZRfAKeoAn3mAT zwuC;eCG_zvp-*TDePT=KlUhQb+!Fefme8lRgg&h$^yw|3&u9sKW=rU^T0)=Q68fB$ zxV1jFC2pzDYY9L8{Fc!3T0&pY68gfH&=<9YzPKgyB`tAteQ8VFRA1H-e*EPvp|5BO zePv7BSYOo=Hw=C0vhUjI-SyRdKjH1KX$gI8OX%xbLSJvEYPxUeG2!}+Eun8}34L=* z=v!Js-`W!To_3Ci>t@>b;oAC(E#d1?Uup^c<(AN2X$k$+me5~o3H|k!(BEhY{mqup z-)af{?UvBrX$k$^meBXMgubsO^!+WNA7}~vU`yzST0%eE68e#r(2ur+eyk<*<1L|| zXbJsfOX#OsLO;k9$HSkw{y+uptEn7mry(RQkEupt=3B65A=xtj1zl zAJ7u|z?RSlwS+#nCG;UJp$}~deOOEA!&^ch(GvQ|me5DFgg&|@^f4`=k8KHkTubQV zTSA}E68glJ&?mKoKDj0IDJ`K-Z3%r^OX$;ELZ8tR`plNlXSIYryCw8FE%Cp47XNf_ zKd-kVyuS-tLSNVt`l6Q57q^7Iq$U1OKZ}2=f8Oi*=Op~|UvHO8CG;&Vp>J&oeNQ__#N;eKIg3Bn#{uE&`u&#BKWGX4!_Zp*Ly?y>UzEO4 z+7f!Rme8BGgr1zmCui}=S^U+{htqpV`20AdCA4!mq0edweRfOeb6VnG>vLP;fA#C) z|I*uE(AyE-|Aj4~FKP*WaZBj`Kb^(>^Zw7y;;v85;*+n7KQaB`J@s|%CwojxzAipF zi!U_&p)mE@w{VXMuN{lDgkH2I^qX5kFV+&ezb1c1W$M|Oem*8=@yS_y=i%dFmmU*7 zUv_N?y<1D@-CIJ>Z3(?cOXxjYLhscQdheFd`?Q4KwO|F_=XrM+Kb@@G^gUl+f(j|0M4_P&~o{q`a#d{hSoPH{%UMro$ zCui~LC&azO*Tg4Z7k~BnSigsa=VybK&>Oad-l!$?#x0>YX$ie)OX$s7Li?IHp|@xW zy=6=2x3`4eswMQ+Eupt*3B7Gg=Nf=lVDx{27q%w}k#dOXwfAg#J-W=pVO)e!eC23oW5vYzh5ROX!zd zLjR;CCVxg{@^$gaS$z6)!E5{EEIv7l?>qhBFmgkSFqTS8yd68hqn(3iBt{Acmcbp5ivpYY=^ZwY-xOXw?GV)DP; zocwP$AME3R@OA6S*TpAi@y||wI81$e&dlR?^|66-r{_(R|BlM!zoW8d9}9$g%UUg= z=d^@gyCw8GEuq(K3B6uR=*d}pau%PQ#b5n=IIf3;&yV98?MbhR|n2qThtrGR8Bv9(|d}= zXC529))Ds=db=cTs2(#&V1K)8m@hJzH75zukl0fo$uQ0 zpNVwszL{$ai6!O)TT6VPYwu_Y{mz!q?`jGC z?v~KEw}k$!me9Z568d*qLf_F6`aLb7f43#{@3n;f{g%*w&=UF&TSEJFnQP*<89jFG zjrLspQ9p0O&;O5GLjOri=s#@<{bwzq-`f)UeJ!Ei-xB%*cB-cPgFPl(|4>Wl54VK= zNK5FCwuJszOXxpu34Lcv=)Y(Q{g*AF|EeYQU$=z5t0nZ`w1ocKme7CK68i3z(0|_& z`X5?Cf4n91KemMarI>Z3+Fk zme8MX3H|Raq5q>L^nbR5{;!tMUuX&KwO~sVJ`V`jCZ7wF&xJSI=fd2nkN(%*7eCzR z!gycQ;~7DZX9VqhKhZp)$NQol?~8i8FY58WsK@)F9`B2Kyf5nUzNp9hq8{&ydb}^{@xG|< z?0F{icwf}xeNm71MLpgZ^D1fcwf|a_jZIH?~8i8FY58WsK@)F9`B2K zyf5lc_V*#Q*MZsB0zKXr^>|;@<9$(&_eDM47xj2w)Z=|okM~7A-WTD1fcwf}+i*roeHlxR`OQ+^Jvv*S;=% zY`)h8J)RNtct+6U8A0FP`z7>vM$qFKL62tyeMfIc=<$r8$1{Q+&j@-vBk1vrpvN)^-C{y5hbtO>nPOX!7LLNC%1deN59Z*B>_SWD>sdgb3y>BVb8FVPbEEiIv! zYzh6=me5PJgkHKO^fE1>mu(5XTubQXTSBkU5_-j!&?~isUb!XoDlMT`Z3(?vOX$^G zLa)&hdd-&5Yqf-4yCwA7rXNm!+Y_&6zt8n`YC^Bu5_-Ls(CfE^-k>G)hAp8tY6-n@ zOXy8nLT}m;(0jFn-n%9AJ}sg5Z3(?!OX&Su zLLbl)`oNaZ2epJgxFz%34K^g=)+q=AJG!}$d=GYwS+#pCG;^Zp^t3|eOycE z<6A|F9+Wk6J?ixFz)SEumj%3H@SA=$BeTzuXe~CoS>8`ll`7`p;TI|GXvM zU;m;d-dF#!CG@XaLjSrY-dj)3&Ypwm*OFyg!fVU2Euoid3B7zv=oMN*uhX{%sKU=qi-liq= zwk@HzYYDx5OXwY1LhslTdZ(7qJGX@1r6u&PEunX73B7wu=(#PS_h<>dXG`e4T0-yL z5_+GO(EGN8-mfL}{w<*oXbF8_OX!1ILLb}``jD2;hqi=1tR?i}EuoKS34LTs=%ZRf zAKeoAn3mATwuC;eCG_zvp-*TDePT=KlUhQb+!Fefme8lRgg&h$^yw|3&uEE<>oZ%z z^|M++pWPDroR-k%wuC;fCG`0%q35-PzMv)ag)N~kY6*RDOXy2lLSNbv`m&bLm$!tz zq9yc|EupV!34L`-=xbU+U)vJ;x|Yz_w}if-CG?Fgp>JvleRE6bTUtWj+7kMnc8-Y2 z*?Dqyey)!L!ddhCEunwV68eWNp?}m8`o}GypKl5MLQCiuTSC9o68hzq&_8L3$=P}G zb>~&4KNO~3+g9x{;k9nHme8xWgkGa1^qMW9*J=s9c1!5DO+TFecJayCd8?iY!t=9r zOXzJ{LT}p=db^g;+qZ(DPbCU(gcz!j{k%wS>O7CG_O%JUKgmclyI&>g%re z_L%T>*!x;S-`^7YftJt@wuFADCG^8Bp&w}p{b)<*$67)^-V*wWme5bOgnp_e^wTY& zpJ@sGY)eeO?)=~Lb>{`9pUzi)k1QYUx%M^b*=twH86KPO+U)D4%a6}@?arCEG3U)( z`_#<$dEc||{|8-LuqO0EEuj}~3B5>5=tWyXzquv!VlAQj>ykdgYeTtF(k(wI%dwEumL$ z3B5*3=rvnHuhkNIPD|*uTSBkX5_;X1(Cf8?UcV*u1}&jCYze(lOX!VTLT}O%defHB zo3(`Aye0G&Eupt;3H|n#&|9^H-nu39HZ7sIZ3(?yOX%%eLhsNLddHT~JGF$~xh3>2 zEunX93B6lO=-pdF&ut04M@#5ETSD*E5_<2J(EGH6-nS+6el4N*ZwY-sOXvezLLbx; z`rww(hqQz~v?cUmEujx@34KIM=p$P~AJr22=$6pOw1hskCG>GEp^t9~eL_p<6I()` z)DrsSme8lPgg&(;^l2@jPj3l*MoZ{3TSA}J68h|x(C4&-KDQYzci+OX!+iM1)Ae&Lp}*e}`Ufqcf7lZG zM=ha$+!Ffvme4P>gnqFl^h+(FUv3Hgla_d<{%K3N{R+_Pll3oKLjS5I z^siguiQ4&lX%n~2=<(fMdq+#?ceaFnS4-%3w}igECG>B#g#PW8(7)3X`i_>+?`aAB zyDg!AuO;;Fw}k$Kme7CL654Yy*TiizdhFU8?Ya1)e%^$i{~x!6{*#u_f7%lI&ssvi zw>wuJt>me6D1fcwf}xeNo@p^GxXRzNp9hq8{&ydb}^{yLvx_9`B2K zyf5nUzNqi+?Fc>I7xj2w)Z=|okM~7A-WT1GBFMdb}^{@xG|X`=TE2 zi+a2->hZp)$NQol?~8i8FY58WsK@)F9`B2Kyf5nUzNp<7=a{%{Mvq;aycWEAM)rKX z-d_vGkA*w?7$E%IC$9@{yw`=fQ@{GJeO>s(e6I_7JR|7wjG)Ieg1)`?OX%^8pvN`6MCVR&8i`Rr+q9ycOT0$?`68fzzp_gh2y>v_HWm-Zn+Y)-Yme9+$ zgkGT~^olK^S855pa!cq{T0*bd5_+|k(5ttEUZW-Snk}K%Y6(53CG^@Yq2D(BaO&G1 z*6Z|$aDCmD(Cf8?UcV*u1}&jCYze(lOX!VTLT}O%defHBo3(`Aye0G&Eupt;3H|@D zcNb80--&(pV|Kh|W{=m*Oft!mLAE3_Gs$e38I#SnP0CHf&@F_UG|;BIh2$KsnVD(L zxM2J3%kX?=?%03pxOq>H+hn_YpQA_L`^;Y&jjm?y{b_bq-l-+@&Ml#LX$ie+OX%HN zLhs%ZdXJXSd$xq$t0nY|me6~*gx;qm^u8^j_iG95`G5Khpl8;EKAbxY`LT0&pj68gH9(AT$wzM&=bjV+;XY6*RFOXyo#Lf_gF`nHzPceE2C{%F)* zhCi&o-4Z^3bZ1NGyIMlu-4gnqmeBXMgubsO^!+WNA7}~vU`yzST0%eE68e#r(2ur+ zeyk<*<1L||XbJsfOX#OsLOZ?Y5`>3xT^*gm+bNbi4pckTS70>5_;K| z(0%*#oUWIv3H_Cp(95@kUZExQiY=j6Y6-n^OXyWvLa*8qdbO6&tG9$+qb2m3Euq(H z3B7hp=yh5`uiFxOy_V4Hw}jrHCG>_Zp*Ly?J<$?+{y+upt zEn7lw)e?H^meAX@gx(-L~$meBjPgx$Fzh#wk7m&EuoKZ z34KCK=o4E)pVSii@Gh0HR)e`#bmeA+4gg&<=^n#Z7 zvi}`mukWH>51}t^34KXR=u2BdFKmfh>dRW<=KAuM@Z+y&34LWt=&M@druyoZxN*?_ zJB~k7U(?G8KmXd6(ATwuzP=^&4R*4o`o?Y(uHV!W`sS9EukN6 z3H?}0=*L?^KhYBU$(GPhwS<1UCG;~bp`UFD{aj1v=UYO*&=UH^me4P?gnqdt^eZi) zUu}tN>hH9~RrPBvpOXc zo@fcZaZBh;T0(Ey5_+?i(3`h}-l8S+mMx*TY6-n{OXzJ{LT}p=db^g;+qZ$Up$}>aeQ-fXInx) z*An{qme4P>gnqFl^h+(FUv3HgN=xWhTVi|`AD_jSnR-{4e5_lx+l0ryzKze~llR8d z{jqMh3HQ%>Euq(M3B5r}=nY#!Z`2Zcq9ydkEulAQ3B74c=*?O}Z{8Ami$HR(e_nk2dGWoc-W?|2ANJ`s z;r(OZmeBjPgxi!s? z#m8sy@mYN81H$9O_$>az`}N}A?a$}D{-}um)~^?LR{l{wi+leapT)=j-_2*IK0GGR z#?N(|a8`B>AD_i1?~d{3#k(68o4P9|&&G>)oA6k&L`&!;TS70@5_;*D(9YrGv-srQ zF?oOZoH*hB@i}oqkI&-cv-tQd{@KUF_$>azdlvt)pBMLd_uum@{^z~E%X-%#{<=Re z{=I(vYx{KxzrXQWe0&xkpT(y>9(;|-yWJ%8do7{gZwdWDOXv?slKt>gzGo8gub~Y^eru+ zZ*2*ETTAFW+6fWkuNNPGz4&Xr9}v#EueXGLqb2m4Eur6P3H^3U=yzH|zuOY}y_V4L zw}k$nCG>|aF+Pir&*H03y(>&Uwyn`^!eiZ_Zp*Ly?J<$?+Kxz2$CKAS z``^nx8U5Jd>zV$1!vFFGujj(4SMwA9W%PQks(*g5iGAMwg4eTP^ke_mFSvI8=-U7G z1=ps3-v9mu*RC1;y#Hf#ZTNdUZ*=Ydn!NVezYqI=F7kT3jw)@>9%V>ZuiAM>i13f{r}^Z(Ep?*^v}11 z{--UWf1xGxFSdmKrIygY+!Fd%T0;M7OXy!~3H|FWp?{+#^l!F={%0+rztIx^zXNX{)3j#|GFjgAGU=4qn6PBrX}mgLrduY*b@3bwS@l9EusIS zCG=mmg!WjlqX{1eglprEh4IJ2kK4z>%*pru&p$6-(#OK^c~K9~2zq!%(9Q>h{(4L3 z;Tb^>&j@;WM$kXo&m;8kjG%{S1U)<>=;0Yb56=jCct+5}#{xZkEYQQ}MLm38)WhdR zJ$zo&!{j@Oe=WpBMG; zc~K9a7xi!TdI&vyUev?qMLm38)W6ryBlPfjQ4gON_3(L751$wH@Oe=WpBMEX_v;Ya zf!UE9zHMX;q#&%J}>Iw^P(O;FY4j*q8>gk>f!UE9zHMX;q#&%J}>Iw z^P={=xW9?pN4o9W__5%#GqU^RNByy2{IUPz_p!g25BB)RTS6~r34LBm=<{1bU(gcz!j{k%wS>O7CG;gNp)YL-y|5+pWi6pEZwY-xOXw?G zLSNMq`s$Wgez~dtPWClZIeGosZWDg|buFQ=9eNRj1ds{-^*An{vme3EhgnqCk^g}J7 zA8rZ#NK5EPTS7n968iC$&`-34ezGO>Q!SyNZVCNNOXz1?LO<6M`uUd7FSLYyu_g3N zEumj-3H?e-=vQ0fSL^Sz#4p#cwS<1XCG;CDq2Fu?{Z>opw_8HL(-QjKmeB9Dgnqvz z^am}WKWvF#slVG2u79s3^!Hohm+Bw1#4pw#wS@k-CG;mP@e8%jjr^I3#YVdAV{7r2 z&`Y$0Ua}?hQZ1pEZVA0iOXy`=LicTacAmU9rtXh*yG^)%)@up9eoN>LT0(Ew5_+SS z&>Ode{_51tsqdE?pPhH^o*>*myR?MfwI%d!EunXB3B5;4=sjCP@6{4|MoZ|uTSD*C z5_;d3(EGK7-oGXE%$Cpxw1l435_)z^=mT3qAJh{1;Fi#bw1l4168g}V(1*2zp4$?7 zUQ6itEujx@34KIM=p$P~AJr22=$6pOw1hskCG>GEp^t9~eL_p<6I()`)DrsSme8lP zgg&(;^l2@jPj3l*MoZ{3TSA}J68h|x(C4&-KDQ-S=<#AfOSpbsOX%}kLSN7l`ofmb z7qx`GxFz%@Euk-M3B9l-^kprfFK-EbMN8-_TS8ye68h?v(ATttzP2UwbuFQX-lYaZj$eobNyT9RBpjw3j}-)aS3w zkRw(a{k)Y&eS1s1-0%0NT0;MHOX#n)g#MY9&|hx}{kK{||LvC0f2Sq%&$fjAxt7p> zw0@E|yr_p~1U)<>Xy*e$f4wF2@Qk2` zX9PVwBj}&)=Mj2%M$p4Ef*zg`^ze+Jhi3#mJR|7gV}Tw%7U<#gq8>gk>f!UE9zHMX z;q#&%J}>Iw^P(O;FY4j*q8>gk>f!UE9zHMX;q#&%J}>HTbUzb%_`Ilx&x?Bayr_rI zi~6^EJ%k=UFY4j*q8>gk>fh_<5qkK%sE5ysdicDkhtG?8_`Ilx&x`tx`*jHIabWtf zKo6f6_3(L751$wH@Oe=WpBMG;c~K9a7xnOYQ4gON_3(L751$wH@Oe=WpBMG;c~N^_ zoG@|wNVi=ZKNfs;Ms|Pvs6Q4A-xt2o`vBqhK7L&I@q1jDIr*pm`NxGf7I|FI!!v>& zo)Pr$jG({X>m~H?jG%{S1U)<>=%4N95qfw=(8Du=9-a~O@Qk2`X9PVwBk187K@ZOe zdic1Yhi3#mJR|7g89@)v2zq!%(8Du=9-a~O@Qk2`X9PVwBk187K@ZOedU!_A-{^iO z^zgAj56=jCct+5}GlKrDUJs#%X9PVwBk187LH}MqkI=(2f*zg`^ze+Jhi3#mJR|7g z89@)v2zq!%(8Du=9-a~O@Qk2`X9PVwBk187K@ZOedU!_A!!v>&o)Pr$jG%{S1U)<> z=;0YbdtTVz#O)*9c5Qe@(BsE~&pwal{`gUQEEs?6e|8S==lkREvA>uP_V~u+_fy~b zy?b2nSO-65`p3DpcunXfT0$?`5_+kY&`Y<3UZy4VvMr(e_UYID=;dlcf2Ae#@-3lP zXbHVyOX!taLa*EsdX<*YtG0w*ttIs8Euq(F3B6`Z=(SoxuiX-QotDt+wuD}34K^g=(#PS=e2~M-xB)pme5DEgg&w*^ieILk8TNlOiSovTS6b#68iX-&?mHn zKCvbANiCsIZV7!#OXyQuLZ8+W`t+92XS9SqvnBLdEuqhD34KmW=yO{_FK7vUUQ6in zTS8yZ68gfH&=<9YzPKgyB`u*ZZ3(@wCG=%2p)YRV z^mQ$vuWt!`LrdrzTSDK|68h$r(6_XNzO^OvZ7reiXeUIxK58$+YxTEV!sn0fYzci= zOX#~>Lf_L8`rek%_qBw+za{hoEukN53H?w@=!aWEKhhHV(U#DUwS<1WCG-<5p`UCC z{Zvcnr&~fl(-QjGme9|&gnqsy^b0MaUu+5eQcLKUTSC9m68hDac(wjcOZeYVUuy~d zdQ0dxT0+0s68f!{&~LYdey1h$yDg#LYYF{+OXv?;LVws2uhidd3D>{Z68ifs@lySR zmUyxLs3r8rEulYYi5F^T=VeSx=G4u}L+f(Ww{5Y)B0cd1J^iv(`tj2r$9~wN|Nr3q zm-N!lzP9MAa>S~Wzwj4dn>tK?{@aUM_;i!LV)K`uA_EBFw>UZkJx=r}^F&A$My+liBzvao__NbR?3D=iy3GL5i`p-r$ zTNAo(pZ=L?zx~l&FnRqeHR1a5EusB+PXC$e6>CDT)Dn8-me8xTgkH5J^lB}kS8oZu zMoZ{5TSBkZ5_;{H(Cf5>UbiLmdM%;XZwb9YOXv+-LT}U(dZH!t#x0>YX$ie)OX$s7 zLT}y@dW)9OTegJWswMQ+Eupt*3B7Gg=NJt0na8me2>b zgg&Sx^uaBm4`~TKrzP~EEujx<2|c$Z^t_hP^IJk6-V*wVme5DGgg&Yz^wBM$k7)^g zY)k0lT0$S+68eOe&?mNpKB*=2$t|HzX$gI5OX$;DLZ99e`iz#)XSRest0nZ=Euqh8 z34Lx$=mjk?{_ik;{_h&@pKH5c2z^~k=<8cT-(V+es&DKz;rdN2p>J*peM?K|TU$ck z))M-Tc0$DUqxLdfSAV-DuC4EE34K>f=(}4&-_sKM-j>k!wS>OECG-O=p&x7s{ZLEj zhg(8F(h~a7me7y2gnqmw^b;+ipKJ;JR7>cmTS7n668hPe(9gAme!eC23oW5vYzh5R zOX!zdLch`y`qh@Wrv6S#`1hV)YYF{&OXxRRLciG(`mL7GZ?}YgrzP~eEur6Q3H^Rc z=nq;#f7lY^|MtuHEWY~GyTasS+Zx>_Jl3t*5_+wc&}+AZUZ*AWx-FsCYYDx6OXv++ zLT}g-dZU)m6D^@PZVA0fOXy8oLT}a*dh?dhTeO7UvL*CZEupt=3B65A=xtj1zl&uj^OKuhRZEum+(gg&q(^g%754{iy4NK5EBEujx>34K^g=(#PS=e2~M z-xB)pme5DEgg&w*^ieILk8TNlOiSovTS6b#68iX-&?mHnKCvbANiCsIZV7!#OXyQu zLZ8+W`t+92XS9SqvnBLdEuqhD34KmW=yO{_FKCG`dlvszukWH>58?G++!FedmiT6U zX-l}iuqAG(FKdaL>&sihkH4ZN^p!24uWE^#>Z@Dg#z9}U>epucOnpr+C;a?tTS8yg z68id<&^Or0n(7<7O}KtjOX!yzM_$)p?i?2TQt}yx7wnn!Jk9BLd zgkGy9^x7?<*J%m8ZcFI(T0*bi5_*G{&>Oad-l!$?L`&$6TS9Nr5_;2?(3`b{-n=FB z7A>K-Yze(pOX#gzLT}R&dfS%J+qHz=z9sYyEunX83B6NG=$%_a@6r-_*Ot(`wS?Zi zCG;LGq4#VFy;n=<87-mrZVA0lOXz)DLhsiSdjFQtGh0F*&=Pu9OX%4xp$}{ceNap2 zgIhu$(h_=3OXx#eLLb%=dTvYTc`c#mw}d{tCG-(3p^t0{eN;>6qgz5B(-QjFme9ww zgg(9{^a(AYPizT&QcLKQTSA}G68hAZ(5JP8KD{OM87-mDYzciS*6lXo{#maj^!hEKH)sjHVN2+Z zT0&2>gxE#H*E>MSxe~6TS9Nq5_-#)&|9^H-nu39HZ7sIZ3(?yOX%%eLhsNL zddHT~JGF$~xh3>2EunX93B6lO=-pdF@6i%^&z8`8wS=D05_<2J(EGH6-nS+6el4N* zZwWoKCG-I;p=Y&(p4}4qz?RSlwS+#nCG;UJq35)OKC~tDVJ)HOwuGM75_*11=)+q= zAJG!}$d=GYwS+#pCG;^Zp^t3|eOycE<6AOXc zo@fcZaZBh;T0(Ey5_+?i(3`h}-l8S+mMx*TY6-n{OXzJ{LT}p=db^g;+qZ$Up$}>aeQ-S^Tx$4+via@_I|?H(El!*%JD#me6mv zgnp+b^t&yg-)jl|eoN>NT0(!=5LOVUShw4R z`)9qD(CfE^-k>G)hAp8tY6(5j5_)_VAAeqa*3`Sh-GujBXO{pEFxRpVboj?3U2yw1hsl zCG>)p__AklukVNHEdD3GJy-SnBF6vU&GEmZ@?`G^gtP7VEI$6c_&2BC9VWk@ek>e+ zz4-XwG#sDB$6uqe-{`Ty*=GN46MAM#=mT0p&uR%hyCwAaEI$4kmH*mV{JXt>oY(sa zG5-HOECG-O=p&x7s{ZLEjhg(8F(h~a7me7y2 zgnqmw^b;+i$7k{J*Qk6l_3kkK|890ad}VZBI2$bAZ9=cm5_-j!&?~isUb!XoDlMT` zZ3(?vOX$^GLa)&hdd-&5Yqf-4yCw8GEuqJs7axC(%HC7&4wLT>`*fS|{;_XMXy@?p zS$y*Dn7luZ=qBO*IkF}6Q7xg5ZV7!%OXy=;LLb)>`uLX6C$xk-ImbrwS<1ZCG-a^ zp+9Vizp1YmUu^1^pL!ccx7Y5m&A+iSymr~>+B#owZTj=+>wdwtt4BX?y)U?S<>=b_ zUvTY`(X|b};MxVFYa4#Swdr4Pqc6BN{p(G9!L_SKuXE!sxHkPdzq-gZuW$NwZZf*| zt1$b9>s3(Cf5>UbiLmdM%;XZwb9Y zOXv+-LT}U(dZH!t#x0>YX$ie)OX$s7LT}y@dW)9OTegJWswMQ+Eupt*3B7Gg=m1wk@81%7W=rS;T0+li2|c?d^nop*4{8a0a7*YzT0+lh34Lfw=)+n<&us}k zuO;;Sme7Z{gg&Arw7=I7+TUvkeRNCcV_HHV+Y;K}YY6S{HH7x}8bY7g68fZ;&?mQq zKBXn}sV$*TYYBaNOXxFNLZ8_Z`mC1FXSaktrzP~cEuj~*#4h!DEwS^cub8ac|8kG| z{B9F|{sk?eFKh{YQA_BHTS8yb68h4X&(tvj)mPSp>sPhJ z4)xV7vHhU^yEVJl*YtA2&%d@M^mQ$vuWt!`LrdrzTSDK|68h$r(6_XNzO^OvZ7rei zXeUH$J8Cb(Hubk#V(a?Ome6;#guc5a^gS)1?`;WvUrXrwTS7n368gcG&=0kQez+y{ zBQ2pHZ3+EYOX$a2LO;zJzts}@?UvB*w1j@QCG>kOq2F%_{Xt9U z4_ji(`nxUR`uAEwf4?O*uYb@Io7Eq+g#NfC^d~K`X}##*XU_Ga*mHDk>8bqmj60cK zo4h=_CRUifZR?3edirD9i~s-r8?N6N{M?6rQhfGzt}6%E_L_VnKD##kxaEh;SmfF_ z20!nS(d+!v(aRQVi39t2i?@Vcq9yc_Euoic3GF_czTfmRHKCVn3GL09e3M4qi$1-7 zUH?i==;d2NyYHv(f4yQ&=#^SRuiO%Pm6p(}wuD}-CG_eoq1R{$y=F`3wOT^2-4c48 zmeA|AgkG;D^!hEKH)sjHVN2+ZT0&2>gxE#H*E>MSxe~6TS9Nq5_-#)&|9^H z-nu39HZ7sIZ3(?yOX%%eLhsNLddHT~JGF$~xh3>2EunX93B6lO=-pdF@6i%^&z8`8 zwS=D05_<2J(EGH6-nS+6el4N*ZwWoKCG-I;p=Y&(p4}4qKs#AeeNeXv*AH$9eMn2_ zIW3_NZ3%r?OX#^Rq35-Pp5GGs@Rra=w1hsgCG=4(p^t6}eN0Q}V_QNW*An{pme41( zgg&t)^hqtDPi_f)N=xWdTSA}K68iL(&}X!SKC>nCSuLT@ZV7!(OXzc3LN90ueO^oG z^IJk+&=UH>me3cqgub{X^d&8!FKr3EuqE_mEuk-O34KLN=qp=7U)2)&>Xy*gw1mF4 zCG>SIp|5WVeM3v=8(Tu()DrsUme9Afgub;U^ldGn?`S7P%pSFuVOIU^mN=lkvnBLh zEursj34KpX=zCj2-`5iQ{+7@Ww1j@JCGoB6=me3!!g#M%@_O6HL?*($iL5p0QeoS@w;6<)oHvN;@L;i|ubN-5Jhc0q$;ox;1 zHoEp_qn9n#62JMApFh*ckTS70>5_;K|(0%*#Os~CplW*ea`d3=Q z_2pYauh0^D#g@=3wS-=|CG;vSp;v7Qy;@7?)muWZ(Gq&ime6apgkHNP^g1n}*KG;C zUQ6ipTS9Np5_-dy&>OXc-nb?7SEp`Hecv6gr(fp!CN-fqZ3(?uOX$s8LT}L$ddrs3 zTeXDVx+U~BEupt<3B6rQ=agg&?>^dT*w z=d^@Av?cUmEurVOgr3(DdVWjj!&^ch(GvQ|me5DFgg&|@^f4`=k8KHkTubQVTSA}E z68glJ&?mKoKDj0IDJ`K-Z3%r^OX$;ELZ8tR`plNlXSIYryCw8FEuqhC3B8~t^m#3z z&uXy*g zw1mF4CG>SIp|5WVeM3v=8(Tu()DrsUme9Afgub;U^ldGn?`S7P{Klxg48LA~yCr_D zzOyCtT`i&SZV7!)OXz!BLf_XC`u>*C5441SuqE_EEukN73H?Y*=to;ZKh_fZ@s`j} zw1j@LCG=A*p`UIE{Y*>fXInx)*An{qme4P>gnqFl^h+(FUv3HgN=xWhTjE#i@3h1( z*RQpNe!V618!e&VYzh5VOX#;-Lch}z`rVe$@3n+}za{hsEulYbiC?L|+Y+vSuO;;N zTjH1MAGE|T)*rQm{VUShw4R`)9qD(CfE^-k>G)hAp8tY6-n@OX#mo-JJSnC zSuLT@ZV7!(OXzc3LN90ueO^oG^IJk+&=UH>me3cqgub{X^d&8!FKr3EuqE_mEuk-O z34KLN=qp=d-g4u!^PSxfgtORPEursj34KpX=zCj2-`5iQ{+7@Ww1j@JCG3yTWQWJXl zmeB6E>HAKvSQC1sme4D=gkGg3^r|hPS8EBqdQ0dvT0*bc5_+wc&}+AZUZ*AWx-FsC zYYDx6OXv++LT}g-dZU)m6D^@PZVA0fOXy8oLT}a*dh?dhTeO7UvL*CZEupt=3B65A z=xtj1zl&uj^OKuhRZEum+(gg&q(^g%754{iy4NK5EBEujx>34K^g z=(#PS=e2~M-xB)pme5DEgg&w*^ieILk8TNlOiSovTS6b#68iX-&?mHnKCvbANiCsI zZV7#govf)owcCX2r?rGWy(RP+EuqhB34K;e=(AfwpVJch+?LP_T0)=K68ij>&=<6X zzOW_qMJ=H(ZV7!!OXy2mLN9CyeOXKB%UeQU(GvQ~me5zVguc2Z^ffJ^uWbo^T}$Ze zTSDK^68grL&^NV&zPTmzEiIvMZ3%r_OXxe=2@xlc+RJcK{q2@GvA(k<^j$5X?`{cw zPfO@~TSDL068ip@&=0hPey}CiLO<3L`tg>~Pqc)7vL*CWEuo)o z3H?k<=x19(Ki3la`IgWxw1j@KCG<-zpV;0=+|39ztIx< z&6dz_wS<1VCGrtZcDiSy_V45Z;4~;AGE|V^+zqC zKW+*ANlP4E56}DRLf^%I`AOg3e%m+Ahy|nl+T^v*9v8np>V>20KQ-#BMt%FJuO9VJ z*Nb(V@NWPt-V%C=me5PKgkGv8^wKS%muU&TY)k0AeR?j}?uyB~V|4v1E#dm|EumLv z3B6)V=#^SRuiO%Pm6p(}wuD}-CG_eoq1R{$y=F`3wOT^2-4c48meA|AgkG;D^!hEK zH)sjHVN2+ZT0&2>gxE#H*E>MSxe~6TS9Nq5_-#)&|9^H-nu39HZ7sIZ3(?y zOX%%eLhsNLddHT~JGF$~xh3>2EunX93B6lO=-pdF@6i%^&z8`8wS=D05_<2J(EGH6 z-nS+6el4N*ZwWoKCG-I;p=Y&(p4}4qz?RSlwS+#nCG;UJq35)OKC~tDVJ)HOwuGM7 z5_*11=)+q=AJG!}$d=GYwS+#pCG;^Zp^t3|eOycE<6A)=(h~aBmKgtcs6KyAzoPr+qwW_%f7}xKla{!= z9-qbC2UCwF%eREbmK9n;uhLT0(Ew5_+SS&=W18H*N{NNlWNWTS9Nv5_G-m)e1 zRxP2oZVA0jOXzJ|LT}d+di$2pJG6w}u_g3QEunXA3B5~8=v`Yv@75A}_mnCSuLT@ZV7!(OXzc3LN92E|JMHw z)oIIp{&VDB-$lJ1LSNhx`jVE=m$rmn*b+bNXYoJn?Y+LYlkoQ6U?=M*e!_nT>&9*q zuHV!W`sS9EWTa}AJ5|BS$x6g!>QdQe14qP652VO&}X!S zKC>mX=bqW;pYQc7?#G|kk0<`QKEEYgzn~@bg)N~kY6(3&i=X_?H-5d!ja|RA???Fl zm$ihxye0G%EpbDAWlLN?X#YLPpRBLy?SvnHbxY`LT0+0oS={yUEIz(2{^-ny%j9+K z$GT07uZxdo@kM4n6ef=)i*}pv*s)kk=*3$?FVPZu$(GRlHGYlC~< zCG;vSp~tiMcoyGf=EGs~@vv*R2_GNxTSD*F5_bZVA0lOXz)D zLhsiSdjFQt2egDfuqE_CEujx?2|d0p{=@saxZj`i`h6m_b2y>DpR@S?&DX_e&)&nc z`6=%l&*Go$_l*C5441SuqE_EEukN73H?Y*Xy@?pEIzqAerR78Uufp1H~IG5>FrZ{ zEb#R!!)q5$ubuW5*G_+vYqRh3j5oP<=Jb7@`6k!So?biaO|H$po^bY?T$}y-o%1Hw zW`EqdZ*lFsH@PIw89@)v z2-^98&>wCIJv<}m;Tb^>&j@;WM$jMW{SbP1M$p4Ef*zg`^ze+Jhi3#md@Rtz#{xZk zUev?qMLm38)WhdRJ$zo&!{f!UE9zHMX;q#&% zJ}>G|^nMBLabWhbKo6f6_3(L751$wH@Oe=WpBMG;c~K9a7xnOYQ4gON_3(L751$wH z@Oe=WpBMG;c~N^_oM+!jj5PEn<(8Du=9-a~O z@Qk2`X9PVwBk187K@T4n^ze+Jhi3#mJR|7g89@)v2zq!%(8Du=9-a~O@Qk2`X9PVw zBk187K@ZOe`nK*@LJuDc^ze+Jhi3#mJR|6j^?nFFJR|7g89@)v2zq!%&>!#p5PEn< z(8Du=9-a~O@Qk2`X9PVwBk187K@ZOedU!_A!!v>&o)Pr$jG%{S1U)<>=;0Yb56=jC zct+5}GlCwT5%lnkpoeD!?RjBW6Sq!v+qLnrV8I!|{qdvzSTKI>FE|JI{eC-q?l0t% zJ#L$PJM+!#G)hAp8tY6-ouozzTk z(rv=^O$6*8i~5|F@cqwi34LBm z=<{1bU(gcz!j{k%wS>O7B{r`wX^GA1OIyPCzpN$ngubpN^z|*FZ)gd9V@v3pT0-C468e^w(08=+L~Jx|Z^MT5r&_}6 zQJ-!J{h5}~pKS^Kxt7qMZwdW{me5~p3H_y(&|hu|eP>JPue5~zYD?(5T0-C568fH& z(D$~4zONv(66?Heyt_etH0h7u79H? z^fy~#o%&lXv3C9KmeAj63H{xcSgZDD^xMspv!*+q#mBSws(maF9&1-?3B7tt=y@%n z*Jug7W=rU`T0*bg5_+AM(CfB@Uauwe`YoY1XbHVxOX!VSLT_v*HPf4Pn{a*8me8BE zgxW$RZHltTS9Ns5_;Q~(A%|y-o7RD4lSW~Yze(nOX!_jLhsTNde@fF z^IJmi))IR6me6~&gx<3y^jme3cqgub{X ze%R0A7xnYKwx1{A=YL>v%V<60YCS68grL&^NV&zPTmzEiIw%Xy=I-&*I}*{OLXp z2xr%4T0%eD68gE8(9gGoexW7wi!GsFY6<;vOXycxLciJ)`n8rA&*J0j;>*l@C`=yP zmhCp-v2MAR(95@kUZExQiY=j6Y6-n^OXyWvLa*8qdbO6&tG9%n*AjY-me6aqgkGy9 z^x7?<*J%m8ZcFI(T0*bi5_*G{&>Oad-l!$?#&%LOy-BwT*T=K?corYe;tM_>PVFY) z^W(IZ(9Yq6KBFb{nJuBuYKcqhvs>bOeO>&Le*F3UIKunCpe6K$Euk-J3H{&cEbiy~ zR%db7$FunOy7;3rAKsJKwIAy?F}^N7p2ZiL`B0cV_AT0N!ehr`Euj~03B5#1=p|c1 zzoR8|e~n+GGPyTq?vL>-KAy#Q8hkwL+-<_=%PuXUcWnthza{i;EunXB3B5;4=sjCP z@6{4|@0QT}w1nQbCG>tRq4#eIeLzd-16x8L)DrsOmeAwt;y=8vi~IdKuiqy^pWhPt z`#FnW){noW_d|?dqcXlOepep{gtP43Eurse34L!%==)ki-`^7YftJt@wuFADCG^8B zp&w}p?HoRy#V2>i__}y^!@@In#pJQkIea{e&)gxNE4(H?zAnDt{#d7*g!^aRmeA|9 zgkHZT^ad@VH*5*LQA_BJ?WAVfYvP3Ko3@1BtR?j3Eupt)3B6@Y=&f2pZ`~4lo0ibq zwuIiUCG_?!p?7Ery<pN~;y>(P zqvF?hZNE-JU)K`)`j*f)w1mE~CG<@#p>J*peM?K|JKA|7#;;Ks&*D$_aX|PQkY`## zKid-ext7q+w}gJ7CG?9epaz6{EB`*&+qq@@cVl~OXv$* zLSNJp`r?+DdltXE>zDR+!uP+dCG_Pjp|5C(@!#DX|J}{|`ZyrGZauy(KAy#IpZRc@ ze0$dP_W18^j{olFs(maFo=aA13B7tt=y@%n*Jug7W=rU`T0*bg5_&v~k7x1mEWY6L z;h1g`K0l6a34L5k=;K>LpU@Ke#Fo$}wS+#oCG;sRp-*iIeOgQC(_2EH(GvR1me8Ji zW}kn)*ME1@k3X*;Ph44_-x99>pFE3?U!yXE|!s^yrCgWzDJ2O9abstB#W@diunn^GCe^>eEVzck(x;fW; zz1>yQA9ur?Yu-F>dhNzJ*JdB1UB2lpuH8Jnw$$|f+}aYmbnVAmLjOcd=nu7o{>hfm zA8rZ#Q!Sx?x+V0_w1ocIme3z*3H@^|p?|(5^e?o8{>7Hiztj@i@8NtCw@!82wYS@S z@yq?X3BUedX$k$SEunv{CG@Yig#L||(7)Ld`nOs_|8`60-)RZ`yDg!AuO;+HTSEVS zOXxpn3H^sHp>Jym{Xbel|4~cm|Jf4yk6S{2tR?iHw1ocCme7CJ68g_uLVvs^^k1}u z{>zrof7KHDuUkU@O-txcw1ocKme7CK68e)Zq5r-m^gpzO{>PTk|I`xt_Lk89+!Fd< zT0;M8OXz=V3H|Raq5q>L^nbR5_E@l)37-dqYvbp___^?Q`&^hm`RIS+^WtuOE)1U+ z_3(_Khi3%sd_d?Aw}c*^5%lnkpoeD!Jv<}mkMw>BJv<}m;Tb^>&j@;WM$p4Ef*w8= z=;3359zHMX;q#&%J}>Iw^P(O;FY4j*q8>gk>f!UE9zHMX;q#&%J}>Iw^P(O;FY4j* zqQ0&BmC(cIMLm38)WhdRJ$zo&AM5=PdicDkhtG?8_`Ilx&x`uwy&pmkpBMG;c~K9a z7xnOYQ4gON^(T72g!VWv`&gid&x?Bayr_rIi+cFHsE5ysdicDkhtG?8_`Ilx&x?Ba zyr_rIi+cFHsE5ysdicDkJul8PaqCpKT^k<@7MzjYA3y4k1;dYp+xi$F{MyIIg}2}1 z!u-ka{u>_`_MG#$poeD!Jv<}m;Tb`Hxc5uw;Tb^>&j@;WM$p4Eg8oSFhtR__f*zg` z^ze+Jhi3#mJR|7g89@)v2zvOqpoeD!Jv<}m;Tb^>&j@;WM$p4Ef*zg`^ze+Jhi3#m zJR|7g89@)v2zq!%(6@EJ5_&j@;WM$p4Eg8q2# zhtR__f*zg`^ze+Jhi3#mJR|7g89@)v2zq!%(8Du=9-a~O@Qk2`X9PVwBk187K@ZOe zdU!_A!!v>&o)Pr$jG%{S1U)<>XwM6~nz(hU+pZ1I2zq=hSnxWU`{PILv0(h%UvLia z`}=nI++WBid)zkpcIF$uc8?1l>)`v$em~b1t_i(JOXx*gLNC@5dhwRfOSFVuvL*C8 zT0-~N>;HWm{jQqO?`{eGo|e$>Z3+Fpme5PJgnoZZ=%rgif1oAwGA*H(Z3(?xOX%fW zLa)#gdc~H|E474Pxh3=}EumL!3B6iN=+#?7&ua<2MoZ`q&fJ`QdrQ4$w+PqQY6-n| zOXzi4La*BrdcBs=>$im7pe6K%EulAR3B7Sk=uKKeZ`u-ivzE}Cw}jrJCG?gpp|@%Y zy>(0IZCXNa+Y)-ameAX`gx;Yg^o}i|cWMc}b4%!5T0(pNpFIQU`8A<;YYDx3OXxjX zLhorOHPd@_n{a*ameBjOgx$Fzh#wk7m&EuoKZ34KCK=o4E)pVSii@Gh0HR)e`#bmeA+4gg&<=^m#3z&uqQq3>@A{Xk3T2U|iv)DrsPme7y1gnqOo^kXfdA8!f$ zL`&!=TS7n867%b?wS<3$`gBX^XIer(+YdXtvWo3@1BtR?j3Eupt) z3B6@Y=&f2pZ`~4lo0ibqwuIiUCG_?!p?7Ery<$Up$}>aeQ-?Yp)YC)eQ`_ZOIku- z+7kM*me7~CgubFB^p!24uWAW>bxY`LT0&pj68gH9(AT$wzM&=bjV+;XY6*RFOXyo# zLf_HO6Y=}g_BMR9{!~l+Uj6Bo(4T1u{n?h#pKA&I`IgXMXbJtrme5~n3H{}k(08_k z{z^;eueOA~t0nZ^Eurse34L!%==)ki-`^7YftJt@wuFADCG^8Bp&w}p{b)<*$67)^ z-V*wWme5bOgnp_eez*QwOZ;~IbW7-GT0%eD68gE8(9gGoexW7wi!GsFY6<;vOXycx zLciJ)`n8t$o%-u7;rcgPLVvR*eyjdgOZ;a2?UvBrX$k$^miUd@nf{X|7Mkj|&#i@9 zLNC%1deN59i?xJaye0G!Euoie3Ef}gS$T4A%-kO!(-QjJmeA+5 zgg(C|^aU-UFKh{YQA_BHTS8yb68h4X(3iD@zPu&$6)mB!Yzci;OX#axLSNGo`r4Mz z*R_Pcz9sYxEun9034K#b=$l(Y-_jELj&`1i@$5XFouBUGfN<7)rX}>VEuo)l3H^La z=oeZ-zt|G`rIyeyw}gJBCG@K;pAJo}-p_#w^>z_yRr{C$@H-D}CXXaeH&3B!$ z)SPRxk3}xOvZQYpyiYazSpBopAGee@#;sG`KB{X!-V*vJT0(!QCG=0Wg#K_#=$~o{ z{nIU>f2Jk$&$fjANK5FSYYF}HEunv*CG;=0g#M+L(C(-CCT^YTwrg*<`{I}TbrXL5 zztR%=S6f2=T1)6(ZwdVyEunw2CG>B#g#PW8(7)3X`gdDG|6WVzkG6#V{g%*w&=UF& zTSDK~68e9%g#M$J(Eqa~^dGl`{#Z-sKWPd5r!ArXtR?iHw}k$9OX$C73H_HXq5rBS z^k27x{+pK2pJ)mFw=JRnt|jy*TSEVROXz=S3H^^Pq5r8R^zAL7|G6dfzqEw@*Ot)# z))M;PTSEUwOX&Y>3GK09GZQ`!2-n8Xh4FLY?e@7afAZ1)#^=Rj`&<}4FY4hLK@ZOe z+WCOcA8rXfJR|7g89@)v2zq!%&>!jj5PEn<(8Du=9-a~O@Qk2`X9PWbEYQQp0zG_Q z)WhdRJ$zo&!{GA@uNhQ4gON_3(L751$wH$9q479zHMX;q#&%J}>Iw^P(O; zFX~V9ehKYyVD_;<51$wH@Oe=WpBMG;c~K9a7xnOYQ4gON_3(L751$wH@Oe=WpBMG; zc~K9a7xnOYQF~sTXX4hWZo4)<7A!a;yFY%^9}9*b3%B(#K=`$fj|*?V$A$Tm-~Bf} zE}SsuaX}By2zq!%(8Du={&4S?(8Du=9-a~O@Qk2`X9WF`-VdRNX9PVwBk187K@ZOe zdU!_A!!v>&o)Pr$aX}By2zq!%(8Du=9-a~O@Qk2`X9PVwBk187K@ZOedU!_A!!v>& zo)Pr$jG%ApekJtqu|N;c2zq!%(8Du={#fsa(8Du=9-a~O@Qk2`X9WH6-VdRNX9PVw zBk187K@ZOedU!_A!!v>&o)Pr$jG%{S1U)<>=;0Yb56=jCct+5}GlCwT5%lnkpoeD! zJv<}m;Tb^>&j@;WM$n!Yb~SP9RJUClo)Pr;Sg_!AH224k+GD}^xxe5X;P?0K@VUQ` zPxiQN^6ktwe(fF?Jl4VYnf-pQEnE|Nk(SVlwuD}+CG_Ggp_gb0y<|)1ceI4=uh+l6 zUcaj*^t)R^zo#Yids{-kuO;+SEur7v5_;*D&>v_Cy-Z8!Wm`fo*Aja9me4D-gkG^F z^hzzES8fTtN=xWfTSBka5_&09im(Gq&gme5Eur^p3B6ZK z=)GG)@6!@`-ciTS6bx68hMd(8sleKE5UN33gI5ePXu>*H3B*eR50aQ(8iw+7kM-me8lSgg&Ds z^qDQ8&uR&Mc1!4UT0)=O68gNB(C4>=zMv)ag)N~kY6*RDOXy2lLSNbv`m&bLm$!tz zq9yc|EupV!34L`-=xbU+U)vJ;x|Yz_w}if-CG?Fgp>JvleRE6bTUtWj(asZb{ItCd z$JL)|39lc0x+U~yT0(!eCG_W7LVvy`^cPw}f3YR>ms&!9xh3?SEup{C68ftxq3>!5 zeRoUfds;%@+Y0(2uu-exfDxlP#g2 zYKdd&ueF4KhWd0%=x16&Kid-ext7q+w}gJ7CG?9epfp zdP}(ejh4{gY>6Z5Z?(h`^|xC>f2Sq%cU$7{`iexQOP~Ga z>vz?Jes@dg_q2q5Z%gR+wS-=(CG`7SLNDDC`U5SYmuU&TY)k0nT0$@15_*M}&?~lt zUa2MY$}OQ+X$ie*OX$^FLa*KudR|NDHCjTi*%ErKme6argkGm5^tvsf*J}yAeoN>L zT0(Ew5_+SS&>Ode-lQe;rY)g2YYDx1OXw|HLT}j;daIVuTepPXrX}>YEupt-3B7$w z=p9-@@7NN0rq35@R-mN9{?k%DBXbHV%OX$5?Lhs!YdY_ii`?iGM zuO;;UEujx+34LHo=!059AKVi9ke1MgwuC;cCG_Dfp^s<@ePm1Mqgp~A-4gnkme9wx zgg&k%^zkjBPiP5!VoT_gT0)=P68e;u(5JS9KCLD6=`EqpXbF90OX#y&;+FdCmbke- zrzL#PuSU#`@Bh@cl1q34M7>yi#A$5;xRW zw#4;=zGUeS?eLTJRlS|?{;qBbeN9W~YgOACG-8eP2uH`&&Xk&=UH=me3Eignqas^dl{yA8iT!SWD=~ zTS7n468g!O&`-6*RrS|e;)?p|me9|%gnqUq^m8qtpKl5MLQCiuTSC9o68hzq(66+F zezhg^Yb|kQ{q>e`{TnTzzu6L(*WYT1%j$2pg#J!T=X^wS->1CG@agg&?>^dT*w4{Zs3SWD=`TS6bv68gxN&_}g|KDs6JF)g8wZ3%r`OX%ZU zLZ8qQ`oxydC$)q=xh3=|Eul|s34K~i=+j$5pV1Qf%$Cq+wZ!*(7I#)YzxP9U{};4` zzOW_qMJ=H(Zi(;rEbjeX(fc9Bv-o%xzpsx2!rAuzme3EhgnqCk^g}J7A8rZ#NK5EP zTS7n968iC$&`-34ezGO>Q!O!`#mBSw5;GqPlgG{_yG?lP>#y;3@yWe0bAPPVZNmMt za!cq{T0*bd5_+|k(5ttEp4Sq3jh4`BwuD}*CG^@Yp~tiMcorYe;tM_>#WedB+R=1gyjpMR5USI%7ev0s>TZT58qmw$22wb>uz+Aqzy z=54drfB8+Wo$EWque`~%+1DOk=>DDiYnc6UzdGl~UE;e=`L#LM&Yr!{{`#D2E}b>K z{u|S4w@-U`ZJUK>-Z=Zdwww0s{j=5N$M!KX{jpn5ug~5Gzu9%aPqV*Q8`r+e^t=6- zUZf`UqAj5pYYDx0OK5+tnf3hB^;$x&-x7L*me3ovgx;tn^u~5l zGrdW-3D-Am3B6fM=*?R~Z_yHZ%a+hvwS?ZfCG<8ep|@=bytRq4#eIeLzd- z16x8L)DrsOme7Z^gg&$-^kFTb4{r&5L`&!+TS6bz68h+t(8sidKDH(FaV?>bZwY-u zOXw3@LZ8$U`s9|-r?iAVwI%dvEul|u34KOO=rdbFpVbmu)@QfG7WFwT;rpN468gNB z(C4>=zMv)ag)N~kY6*RDOKe_W(h{51m$roOe_2cD%UeQU(Gr{1SGL3^gZ4E)+tgR} zcEXRpx+U~AEupV%34L8l=<8cT-_R2J#+J}GwS>O8CG;&Vq3>wtiP&h`-i8h9PqoAb z^`~1xf2Jk$XInylt|j#6TS9-KCG;0tLVu|x^p{&g-`NuSD=neF+7kM%me6;%gubUG z^t~;i?`sKte@o~GT0%eA68fQ*&=0qSexxPzqb;EyYYF{$OXw$BLO8hd;*)!0=Kff@+l2dPm6p(}wuD}-CG_eoq35-PUZW-S znk}K%Y6-n|OXzi4La*BrdcBs=>$im7pe6K%EulAR3B9qM)J$*EZNl|UTS9Nv5_G-m)e1RxP2oZVA0jOXzJ|LT}d+di$2pJG6w}u_g3QEunXA3B5~8=v`Yv&u;(0jFn-n%9AJ}sg5Z3(?!OX&SuLLbl)`oNaZ2epJgxFz% z34K^g=)+q=AJG!}$d=GYwS+#pCG;^Zp^t3|eOycE<6AN3qxQP^?{t5RXYuiM@q1@JJtk+_@hm>RF8 zlfl_xd|iBeU3`39e8K0#cozTf`!y=R+vC}F{Zz>7eM3v=8(Tu()DrsUme9Af zgubJlCt~~>mGNs-p6=s-@HHLJw1j@PCG>MGp`ULF{X$FV7h6KV)DrsTme8-XgnqRp z^lL3KevQibH7d)@d?-vF+m`J%;jwPHme9+$gkGT~^olK^S856E96p}KCwIr>{us~V z<5_$>i!bb>E}iG`Ci)+ zdOVAdUoZa1%!k9|b@)fSP52s<$67)^-V*wWme5bOgnp_e#;;KszeZ(=nGc1@W6zS^ zCOj6sqa}2IjbEcOxi@C+kMS%%e!cikgO7)uyG{6f*`+1)t}UVGw}jrUCG_qsq4#JB zy=P14y;?%=-4c4AmeBjQgx;?u^!_cO4`>N}U`uG{@bN4@xjQEJ$I0C!+&`zZgg&(; zw69Mg^yw|3&u9sKW=rU^TH>Si*)8$Ce!aMVxAOdc9O3={D4oUqdd9Q(`0s8$H1pv- zIU7IRZNgdEIea{ePwtNKb@A?og=X%G$z$Qd-6lMiEYcEs(U#DQwS->0CA4$+cov`B z9h3XRYvP3a$7|w*_L?}M$Ful&79Y>z3qBvlv-tS+;{Vm_;+OS!c}X8C#Q*F6j>_-% z_rJQoFX8cfJd2NK@$oD^^ZDRwOkV0Hp_Zp*Ly?y|JCtOmEU{!u3sCLT}a*dh?dhTeO7UvL*CZ zEupt=3B65A=xtj6qgz5B(-QjFme9wwgg(9{^a(AYPizT&QcLI`JT0&pg68fT+&=hh{!J$FEW8 z^WmlL3&Pp_<(AN|w1j@OCG=}8@q>C@eD>=|{2XTgT)pTb>egFkzotaZKF%%PwGVnj z{MthEcAEJj>^J^@md={F{9{W@-_P*xanAJGl5cWt_Q(BE_l>{*@W)*>{c-P{^W(hz zwCS}!p8mMSr~QuUnj5<|xjcQFe$Sli^XByI>lohrCw>1r2G>6}d1niLj>GT!r*p2& z-d}$BpG~hVF?e6|rnf)dwc*eCmU{Btr`PrH=d5RctnZ>{U(3*U_RfYs=Rcpkvs-Uj z(#0v4PWR@1Lc_oJt$p9gcb{I@!@swlH@#-`d4s;8_vg9nqUqcIvg^a2%U{&P-{aY@ zMfKLdnsaRjZ=CYiZ*lE!-s0LP=3KjS_8qo=JH7VF>DzqV5r5E~vT#diAA*w)!)Y(t z60R@S657Y(?8l_`dq4a8ul;_^{+{Umdj0o8`{_-7imiTEOXzpEg!cP1`+KF|TNCK-Yze(pOX#gzLT}R&dfS%J+qHz=z9qEJ|Jla? zy<<)2omxWg+!A`1me9Mlgr46LdbgI)ySIeiqb2m7Eur^n3B7kq=zUs3@7of3zn0MZ zw}d{RCG>$Up$}>aeQ-?Yp)YC)eQ`_ZOIku-+7kM*me7~CgubFB{_dS`{Qp+DepPQL^wlk) zuW1Q=ZA<9uT0&pn68eUg&^NY(zNsbj%`KsCX$gHtJ5R*#Putt@(fU&@@q6{BTS9-P zCG=-oLVvC$^ygbbf1xGx7h6JqsU`H6TSDL268bAGp}*P^`mUDHcejMTrzP~gEursg z34MP{=m%OtKiCrbp_b4Ow}gJACG?{$p&x4r{di00Ct5;3*%JDxmiXQJYc28H_0uh( zpJ@sGY)k0pT0%eH68eRf&@Z-xeyJt&%PpZ_X$k#mOX$~H;&f2Sq%cU$5&>ft&5gL29z-|5DJKOpq_T+Cd zz3m`8xjg9i%-*(LoHyuytY;q!y!pg>_H$eBx~Oj#{2b4dQ~qi4H(GFQ`18Mg&NV;e z-06?|=Q-Dg@AKE@T$}wde(1lHvT__-}q-W zuhA~u{Y>Z&w1i%!CG@f_p_gk3y?jgP6t zmeBKCLhsfRdiR#ld$fe!vnBLiEur^r3B6BC=zUv4@7EG~|CZ1Pw1hseCG0))M;gme5DEgg&w*^ieILk8TNlOiSovTS6b#68iX-&?mHnKCvbANiCsI zZV7!#OXyQuLZ8+W`t+92XS9SqvnBLdE%C$l^>;t$e;3C6b9MI%p|5EPeQitV>smrz z-xB(Ume4o0gubaI^vx}yZ)pjAM>|i%`1N;h{Ii;WPyF8Q7ee3H68ip@&=0hPey}C< zLoJ~nZVCNJOXx>iLO<3L`tg>~Pqc)7vL*CWEiwK#qTcvtwekNcyx^akcx-q_A0veB zuh;*qsNYo+`rR#|-_sKMy)B{N*AjZEmeB8S3B7bn=nu4nUZy4VvMr&PYYDx4OXw9^ zLa*2odZm`oE4PGRr6u&LEumLy3B7tt=y@%n*Jug7W=rU`T0*bg5_&v~k7x1mEWY6L z;h1g`K0l6a34L5k=;K>LpU@Ke#Fo$}wS+#oCG;sRp-*iIeOgQCALLp5d;RaVUeV9@ zqJG|lpZ~=zF@F8scozR^9|wf9=v^(L?`{cwPfO_WEWY5Mo2=XY@oe`Cp`U9B{d`O4 z7g|EU*b@4sme4P^gnp$Z^s6nQUu%gU)YsqLKJ&Mqe0$dP_V@d``1tR_jA!xjEIz(2 zzWem?#Mx$#ZWDUXme6~(gmw-e&*GE2V{(7IozCL_U%&rn_e&?f*Z*FtkDv4VafH60 zCG>?Yp~tiM_`3L=Gan9 z@pbW;&jrt|8+4P<8@7brs3r8qc2YCFNw*2tH*E>MSxe~6TS9Nq5_-#)&|9^H-nu39 zHZ7sIZ3(?yOX%%eLhsNLddHT~JGF$~xh3>2EunX92|d3h^lmMoox{ho_~h=G+#g4D zlW_kW*%JDwme5DHgnm1n#edkoM&*ALz{h=MU>G$`_exHdu>Z@A9@Bh^;p|5EP zeQitV>smrz-xB(Ume4o0gubaI^vx}yZ)pjAM>|i%_%$ly>*7!MaX@(8`I(l`4Yqds{-kuO;+SEur7v5_;*D&>v_Cy-Z8!Wm`fo z*Aja9me4D-gkG^F^hzzES8fTtN=xWfTSBka5_wS-=~CGgIhu$(h~a6me7Z_gg(3_ z^mrEk-~W2?Pj!EduZxdo@z2hD`b=Jz{#>^SuTy`%CG;0sLVvL(^p{#ff4L>}oh_li z(h~ZsEursf34M1<=zCg1-`f)UzLwDUw}gJ6CG>+Wp&x1q{cua@M_NKZ+7kM)me7y4 zgnptW^ph>2pK6Km>&1U~Ul$+$-OceVKAy$L*Tu)z#b=(Y#*C)3CjI@cY&rWn`BmK}{Cuu%34KjV=xbX-U)K`)`j*f)w1mE~CG<@# zp>J*peM?K|@pbX>EdI5b4~NO?-s9`y3;sJhzGme+-8Y2(ZcB`>i;w^A=E^f43X{ju zRk}@hY+bb_^lB}kS8oYDuO;+&79Y>z<5_&c=fe@*Bz%4x*%JDwme5DHgg&Mv^sz0W zk825id`svPT0)=L68fZ;&?mQqKBXn}sV$*TYYBaNOXxFNLZ8_Z+CP(>{pYejYG0#r zVZXjB`t=dt@Bhcl{WYG&$JfOln)z^=yiWaaw+XLXKhhHV(U#DUwS<1WCG-<5p`UCC z{ZvbguZ#cizAirgyPM*C|rsLVWm zIQu@>O+q_|k7x18-7$F`{-tgbz8?MMme8-XgnqRp^lL3K{&!Rs{J)?)cP-X^Lg>X? zLNC!0ddZg1?`R3#U*mt%aB^?V+#f4rme8xUgkHTR^t_hP<5_$>i;rjV z1)mS&S^R(hEdIT|F8=8rZ_ewkCOrO}-xB(Qme3crgdV>}W&9eIJ7+!|CTHF8>&3^{ z#lJrNcyLz#Mz;z5&6fB-_1{rhXy)BbzMcJNgg?KCT=4%=QdjyWe&H>yeX*b8fA!;j zY0i(kyzfceIp@dSHhJsJH~!zh!}mG+bM)4)yv4PzPOr_L*Jsb=ew_cW@8q%w{hpbd zulKw;J^MKC&Hvgv^|gc^zj1dz{*IZOlW*^uyz>S3-))0y{=dI|Jqr)scIouC|F7?5 zgSX%F7VqcQIoI4Tvp>gs=Ukh8|K9a|bFR(v#wqu|#kB|C;@X3Au3a+w4%>&O*FHIY zn~%5Ik2(K8Klj7zeW87rOg>avy=Y5le?BAh@Xu)4@AvHQxgP!*P50O9e@4@OdXt~x z^!>f7C0u`ZOK2Y_vmYz^y)~iV*AjZEmeB8S3B7bn=nu4nUZy4VvMr&PYYDx4OXw9^ zLa*2odZm`o!#|_x;h)j;@Xu&^_-8ad{4<&!{+UdFaOQ6{^KEy$o_(9^Yu1Ebt0nZ> zEuq(G3B7Jh==EAcuip}SgO<=6wuIiOCG^HEp*Lv>y=hD6&00cl-V%C?me5TEuoKV34MG^=o4B(pV$)mq?XVpw}d{WCG@E+p-*cGeR@mi zGg?BQ*%JD!me6Ongg&Pw^tmme&ua;NeoN>JT0&pg68fT+&=sPgezPcs!H7%j9Z3%r{OX%xcLf_C5`o@;fH?@Smxh3>1Eursd z=ZP5qf4y(~cY6K*ZSU=VA@qGMq3>@A{Xk3T2U|iv)DrsPme7y1gnqOo^kXfdA8!f$ zL`&!=TS7n86661`_l;k>^!@(-das++)M=&qcG3A>R7PoNL}T`xyG@^xBf% z82y$lplg%Mvu`QhGke?C^XByIW3M+q*7e!PI{kP(`&gk*tUVsgGWY-IJNr20`V)QM z;jeG@G0j__oO5mV@zS-Y=3E=zFJGH;ZFqk^J?Gka-r1CA=3JY7e0TZTIoD>Nn_N3* z&b8UcY1iid|9rh|_V2m)B5M5aD!CWDPWtYa&^{(-KQ8rqYeK)TCG=7)q2J#Udg+$X zA7}}^OiSox|9|%G0$Q?zz}G#VgdsxQ(@n?ST{`aW?twUv0P!%w7~}4f5ylu-MtLuT zTikU*TxJ;H>Am;e_xLT0$>q3B9l-^rDu~8@7brs3r8qEulAQ3B74c z=*?O}Z{8Ami2 zEunX93B6lO=-pdF@6i%^&z8`8wS?ZgCGGEp^t9~eL_p<6I()` z)DrsSme8lPgg&(;^l2@jPj3l*MoZ{3TSA}J68h|x(C4&-KDQ-V~@S3+Ok68eUg&^NY(zNsbj%`KsCZ5N0b|Blm_{~pY9%x4eQ!(X`&vTZ-xB(Pme3Ejgnpf3NBNB)tFEwuJ8`zpf=5zrH2(4K1N>Yzci+OX!Zk-4gnlme9|(gnq6i^z$vDUuX&aVoT_kT4MbEktNSf^11!S zt`kDPsU`F>Euoif3B6oP=>GfKpU3F$sR{j-me9+$gkGT~^olK^S856U)|SvKw}f7$ zCG@H-p;v1Oy?RUNHCjSD4BU?fr)e`#X zme9wvgg&+<^l>erk8cTmLQCiqTSA}I68hwp(5JM7KD8zEX)U2oZwY-yOXxFOLZ8(V z`s|j_=d^@AwEuqhE34K9J=nGrom7a@x`HOow;+y?-@uz$JSN3{|ZR>Zogtz~$ zmeB8R34K*djDN4vYhLrQt_Q-`SwG$q`V%doKiLxcQ!SxC-4gmUEulZ#68dv3p+Dae z`U@?gZ)*vCdrRm$T0-C168f%|(08|lzNaPhy)B{dYYBaSOXvq$LO<9N`k|K454VJV zq$TvDEukN43H^9Wj9(WYzb?Ma+)ZI}U$<vWlL{jA#( zdcBs=>$im7pe6Kzme31ZLN96wyn;tme9Mkgx0))M;g zme5DEgg&w*^ieILk8TNlOiSovTS6b#68iX-&?mHnKCvbANiCsIZV7!#OXyQuLZ8+W z`t+92XS9SqvnBLdEuqhD34KmW=yO{_pVt!l{Fcxcw1mE}CBD&faj)-^UJs!!Z3%r@ zOX$m6LSNAm-^#Cx&;EUZt9v`dGrj%yw1oHPy)B{N*An`gmeAL>gubpN^z|*FZ)gd9 zV@v3pT0-C468hG5fr#-rD&x8M6WtC7Uw8guOXw$CLO<0K`stR?&$NVowk7m)Euo)p z3H?G#=oec;ztj@nuCI%~cJ9qjetgID^8RmdVafL*tJ(e1v%TN*kIn8=4_L~}UF3g{ z1DA4acHeW*QjX2umuHsq(l38sW-oWhQeJL$?9i`s?8SM%hS_6;xAfAyWAFBpQ~rJ4 zv0G*bTHg;od_O&|T;Ig0Z1z~^<-R=c*xUT%l&^f9 zV_*F`$Nsl@$A-W5|NZM6`#(oar#`@Cbb$9{+Z$JaSF|NAaxkAYt9|C#r4K7KR)=e%RX>t*f_Wxh4v zuQC4Hd*h!=E$l`}I9FQK5_-dy&>OXc-nb?7CM}^iZ3(?uOX$s8LT}L$ddrs3TiHp? z_10Y`9QXTRLXW=}9)B-9{$6;==Yw+t!so|%Euo#`%+7Q41vQ~P7S0|IzuSH<{LB6N zuIkrEjDIdQ{<+lM-3|zU&T~&o=zCj2-`5iQ{+7@Ww1j@JCGO1$(rN(pd@mzd77a#vzYQMR=>dCp*{#_=VdmYdc`oNaZ2epJgxFz%< zEujx>34K^g=)+q=AJG!}$d=GYwS@L(QiLA=TykcBgpb#yEuk-K34M7> z=qp;{TlwcwUjO*-V2{rc8_&hZKbQK5JA)~+uS@Lwy1AdYX!`TJPJ4FWv}5;GnEgI~ zr|I$8pAU@Z;^VpacrHGki;w5xbN7?Z?L0q>&?mHnKCvbANiCsIZV7!#OXyQuLZ8+W z`t+92XS9SqvnBLdEuqhD34KmW=x^uO#joh?U)=jacz-Tw3H{xCE`E8xzH9n*5`Mkc zw#4|o)8qGRj9(Z3uk{?U*`FEyw>S6|e|t3P*VlyOJ`|G=$F$$r5{|#ACG;{ap_gq5yEuq(G z3B7Jh==EAcuip}SgO<<>T0$>u3B9N#^oA{=H);vJaZBh;T0(Ey5_+?i(3`h}-l8S+ zmMx*TY6-n{OXzJ{LT}p=db^g;+qZ|D_dgi`kgJ||M$Mmy}xt)?wWA? zs+L%*zPcsW9JKG9d0YLSexC61?`;YFzLwC}w1mF4CG>SIp|5WVeM3v=8(Tu()DrsU zme9Ah3q-6wZ9j+A>VIvCRqKzng#LI-=ufnS{$xw&Pql>pbW7;Zw1oa_OX$zFg#LU> z=r6Q{zO5zn?Jc42XbF91OX#~=Lf_pI`kt21_qK$-uO;;TEukN13H@M8=!aTDKim@f zk(SVpwuFAHCG_Jhu}b}KE%Da+iI&h`Yzh5jOX#OsLOcHXaH`1ZD&cg*AI?Bl!ryko=Hxx>6;SNC#PP2b**OF1?@^qQUK9UH#Ro#!2! zJ+^q+U8cuYn!cj_bO+OCVA>eo&9_5XBB=$~l`{j)8hf379; z&$opBg_h92*b@4eT0;MFOXy!|3H_@rp?|F<^sl#s{*9Kzrof7KHDuUkU@O-ty%Z3+E%EulZs68i63LjOZc=znYp{ZB2S|G6dfzqEw@XiI4K z1>2hNc|bTeelCok3$M1%g~gLw|H~g2{eEHixTuHc2zq#qpq(EO`U5SYhvx`-c#fcl z=LmXuj-ZF<2>OG)UP2Gg5%lmJK@ZOn^za-(5AO@~@V-C~9~brTaZwK+7xnOQQ4b#% z_3&{~4<8rx@NrQO9~brTaZwK+7xnOQQ4b#%_3&{~-_rF+=;7m{9zHJW;p3tnJ}&AH z^?C?Bd|cGS$3;DST-3wIMLm36)F1Bk5_dic1ghmVW;&wD+Dc0Vw?FVMrs zMLm36)WgR`J$zi$!^cHEd|cGS$3;DST-3wIMLm36)WgR`J$zi$!^cHEd|cEX7Z;d# z|5TS98}AF2oFltFzT58$hPQ=Vx(yJ1?c@EztM7hc@#Lfb^7{o}Lmu8Q=;1kn9-brU z;W>i-K(Ck3!*c{ZJV(&Oa|As+N6^D_1pUEYFQJF$2zq#qpoixOdU%eYhvx`-c#fcl z_X~P>j-ZF<2zq#qpoixOdU%eYhvx`-c#fcl=LmXuj-ZF<2zq#qpoixOdU%eYZ|Qm@ z^zgnw56=Ly z!*c{ZJV(&Oa|As+N6^D_1U)=Q(8F^CJv>Ly!*c{ZJV(&Oa|As+N6^D_1U)=Q(8F^C zJv>Ly!*c}fabX`5@1N?jW8;0nl5+&t$9MaE!T7nqSL-nS+6el4N*ZwY-sOXvezLLbx; z`rww(hqQz~v?cUmEujx@34KIM=p$P~AJr22=$6pOw1hskCG>GEp^t9~eL_p<6I()` z)DrsSme8lPgg&(;^l2@jPj3l*MoZ{3TSA}J68h|x(C4&-KDQqQq3>@A{Xk3T2U|iv)DrsPme7y1gnqOo^kXfdA8(2A zd1>FA=cRpQ?&_TUIJ++#&&9`c@%6ez5FVS>Zwb9YOXvkHp%=D{?&f5{_Tr68eUg&^NY( zzNsbj%`KsCZ5N0bzei*|7k{GL0pa@~POI1eAs#pkXN_ZLTZkerk8cTmLQCiqTSA}I68hD2F8-bN9F@Jg{;%xgK#YHP()e}p+q)eQ&Smdt34Lcv z=(}1%-`x`Wo|e$}wuHW~CG`C*p&w`o{a{OI=i%eI_~hyszb@X@@VdFHV)EGHJbXMC zpSwal)~?z`La){mdi9pjYuHK6_4sx1@%iGr4Q>y+cbRZ|*`p=&o-LuBhmYstldEHL zeH_(A!u2zri~rYozPNMc@$2H_*Tp|G_vt%17yWFP3FoGtYYF}Nme5~l34L2j=M`e7D%J>|W@$2GqcZlP;_*C|{#lJdtcks1;Q(!z7pS!OdpD#W> zUwo_S=Z445t-DO^~XbF91OX#~=Lf_pI`kt21_qK$-uO;;TEukN1 z3H@M8=!aTDKim@fk(SVpwuFAHCG_JhG5$R&-<-ckW%j%g*XiuOaKEl=Ts$|g*>}=Q z{_U#avHhpVhOfu>)!=~X>v`W?zUIKM^Lh@NckJ3;X7=_EUdr2;9XoXTHm;n$jYH9ND+9a}K{d2{%@W3%hRu_NXkdz+t}a^$>Ym(32eN6kBS)!^9C^NtPQ-ZArz z4PWO+-|Uqv`B+>%c%AeAy|dxl+hFo!`=!6!f_cZ@J9wSH@Wz+Fz1icZpFMHj`#F28 zcI>2i$Go4j_vhrL9P=_bW!^C#=LHVLsq>EQKKIj;A5WWiY1jNa~Rd%c8n4xcBp&lUapn$W&yg3zx1$%kOtWouip}SgO<<>T0$>u3B9N#^oA{=H);vJaZ70D;e_6_CG=)3p*L>{ zy+uptEn7l64=422Eupt*3B7Gg=7fr`v-2qSv*Aetk>mH?)L) zV@v2awS;yaKAwwDu8zs|v1S(u*Uwrlq1SE+y-rK$@mzd77az~XmwY}P)kVVR$I&gJ zk7)^gY)k0lT0$S+68eOe&?mNpKB*=2$t|HzX$gI5OX$;DLZ99e`iz#)XSRfP9{%k* z7ym||tL6RwR-TK$(&uWO-TO11i@)YI|2nyK&7B*5tjmOR#gDgy{zOaYPqu{qR7>bj zw}k#oOX$zGg#KJh=+C!={z6OW+gd{3-V*wbme6;$gubgK^mr~le&57b=WY(;_f2#? z^eHjHefL|sMCj#PLa)#gdc~H|E474vYfI>rTSBkW5_;8^(5tnCUcDvs8g^21J)Vn? zUl-qPaC_Lj%Y@s@9xb7LO`Ooq!^d;+$<;BrK91@l;rcncCG;^Zp^t3|eOycE<6AaUBR)7!hYw?laQ*R_Pcz9sYxEun9034K#b=$l(Y z-`Xw^G5(uZVEuo)l3H^La=oeZ-zt|G` zrIr}~&8zY2;^Vpa+~(0IZCXNa z+Y)-ameAX`gx;Yg^o}i|cWMc}b4%!5T0-yI654tAcrHGgoT#pnh$a3HJ@2pF(KQPa*UhT0+0ECG?wGLNC)2 zdfArH%e91lb4%#{n>g#Ybb)Yu`IgWtw1i%SLvn`=N*An{kEup{A68g55 z(6_gQzN01doh_m7Y6*RJOXz!ALf_jG`o5OX_qT+8pe3~P@bO%Ha&?Sf7w>8q&&B6% z3m$7$?joUAX$ie*OX$^FLa*KudJQ|Nxn8r&gyUbZVA0l zOX$Tdq4#YGyagg&?>^ms1*-Fq(no%S4+y}G?!+3k*arOy|4 zd%dccBVOrqRL<|^-`mR(Uf=s#LSNGo`r4Mz*R_Pcz9sYxEun9034K#b=$l(Y-`Xw^ zF+N9SJQshW+X3M@AYW_={bWn%r&>Zk-4gnlme9|(gnq6i^z$vDUuX&aVoT_kT4H>T z%J_Bh@mzfFbHRQ4#$6=zCM}_zhZA}{7az~X51qR?Ol}W{b(wJcIJ_nF5iOyQYzci- zOX#CpLLbu-`q-Ax$F+n$z9sYtEul|r34Kya=#yJQpVAWg)RxevwS+#sCG;6Bq0ejy zeO628vs*%+(-QjJmeA+5gg(C|^aU-UFKmhLv~%(AwCAW?(DgN*i;v&CdH>v}%lLKi zJ|D(&@wpqr_#Bn-IVuafDG(l87q*06)Dn8bme3ougdWeu$FGYYG+9mLoqO|>AFr8SzNq)@HT}Cv zubZ}+eNE%-lVeN%yzCv*|9AL$X7^ziFXi>j?k_Kyz8-&W_p`5Eu;=MZc|EhQlU%x# zV{f0nZ9OIj@uKOmcTSEi`TJq7ocDUXj&rB~@4Mz5d)xHq z&AY$Ov8$%XmKps2@0?y!O#15wcb3hm!H=Q=Jo7)_S0+k`feQ@zk1&LHT(LM z|K&aNj?J!1$KE^d*zDu%*xM)X;L`t1t>1Z*|6|Fq;p?3LZ)(jxu3qAI=e^wQW9Qh1 z<{fi8nsM#C_jC4V%Z^_+@0japfdg^CbQH9h?1mvEw&Rk4@gCdib$Btz*NF zogRMd^a8JBiXPsEZmx$PyPN9abvgWhe_(QBdim$=?0n_@lb2ia=i0M>f9@l*H{8xX zF5l?)i;Vx)SNHE-pM-O#m0Lou(h}OAiH(0QHn}<`*T=?PBwRn6w1nQYCG=)3p*L>{ zy+uptEn7lw)e?H^meAX@g!cS2LT}d+di$2pJG6w}u_g3QEunXA3B5~8=v`Yv@75A} z_m z34K^g=)+q=AJG!}$d=GYwS@L(Vubc*VuU`nCG>GEp^t9~eL_p<6I(+2GciJ++!Fef zmeAv$i+$^UFaGU*j+&3}d;2&O`h6{-uW1Q=ZA<9uT0&pn68eUg&^NY(zNsbj%`KsC zZ5N0bpQAQDN9~Dj2ZZ12ez7I=lP#g2Y6<;xOXz1>LOymI^wQ@`7Ra!!?+7f!Tme8xWgkHl=YOdGp zGU51IEuq(L3B68B=yh8{uh$ZK{g%)hw1i&J5_(}v=tV7|H*5*LQA_BJTS9Nr5_;2? z(3`b{-n=FB7A>K-Yze(pOX%@jd^{H)&&8K~KAhb}!so|1EuqhC34LBm=<{1bU(gcT zW6$jI=No;F+Tvc{CA~hv>%X)m^kprvPknhyIDSP-jDIdRo{Qhs?SSyT(YLpRzN01d zoh_m7Y6*RJOXz!ALf_jG`o5OX_qT+8pe6K!EukN33H@+O=to*YKiU%dv6j$}x5W7O zD2;!Q(lT>5g~@&AvRx+JmoC>5`pqq&`)}f`-_ixb@#R}Wuh0^D#g@=3wS*qe#m95; z?dNU|liR}%T_)TZgg&q(^g%754{iy4NK5ELTS6b!68i9#&_}d{KC&hBQ7xg5 zZV7!%OXy=;LLb)>`uLX6C$xk zq0ep!eNId0Z|1r9_xAcO>h%(@^zTvf{#??_5&F`W(3iD@zPu&$6)iEIi;w5xw{<%p zoQvMx654qt|j#LEunX43B6-W=$%?Z@7xl4mzL1GwuIiTCG_qsq4#JBy=P14y;?%=-4c4A zme7k^LXYR-$im7pe6Kzme31ZLN96wy+}0dlw1U&mJwI_iPEhS4-%0))M;gme5DEgg&w*^ieILk8TNlOiSov zTS6b#68iX-&?mHnKCvbANiCsIZV7!#OXyQuLZ8+W`t+92XS9SqvnBLdEuqhD34KmW z=yO{_pVt!l{Fcxcw1mE}CBDDDs3pd)i;w5xpX_!(_`2+;T0(!iCG=-nLVva;^ygYa zf4(L37g|E!))M;mme6;!gub&S^j$5X?`{cwPfO@~TSDL068ip@&=0hPey}CiLO<3L`tg<+zb^ir`*rd0dpG;NR<}{YW6jDfp;u`My=qJ7)mlQY-V%BZ zJE^%|v&)3zYqf-4yCw8GEuq(K3B6uR==EDdZ_pBYK}+a`Euj~+gx;_v^hPbAH*N{N zNlWNWTS9Nv5_G-m)e1RxP2oZVA0jOXzJ|LT}d+di$2pJG6w}u_g3QEunXA z3B5~8=v`Yv@75A}_m34K^g=)+q=AJG!}$d=GYwS+#pCG;^Zp^t3|eOycE<6A5 zq3>%6eSb^n2UK2y;+D{tw1mF2CG=%2 zp}&>q;=A_xuj=g(uk_#g^8UTImm~E1T0&pb68hSf(ATwuzP=^&4K1N>Yzci+OX!JcwuFAJCG_(xp`BPl~w0%3X}V`)w)c$uUow(^cr?jbG>Gl3CGuJ3B7hp=yh5`uiFxOy_V4H zw}jrHCG>)p(9XjNJ)Vn?Ul-qh?&dJLJsi+w!tLY0me9__$8+(?)iJp~PVOS%`Z=W~ z^r?YabA5)Ar6u%NTVjRU-&6h(Gi7#PxNvf8$$jCXY0v)5Zi8vhexE=7TVLbf zqq3me5aF@ac{rgLwS?ZVCGxOLpU@Ke#Fo$}wS;~(or}NH=cxQZ*U!aWPlW61l9tey zwuHW{CG_Pjp|5C(@i{8v*Trw^c0f2czr7{&9W9~nYzci=OX#~>Lf_L8`rek%_qBw+ zza{hoEukN53H?w@=!aWEKhhHV(U#DUwS<1WCC2Bdd~?plUpseoPJX<4dieug*Ta8L zNL}fN_`#(dyKwH%Ykp`c$7ZkdhnI5f-P4!*k)<4)-B14LQjWc2`f?v!%CXto`>~}Q zo4vgsU&^tormyoSmU3+NI)8HBF|Tj-I)7?|^yS_3LUvzrH2(8(KoYu_g4IT0$?=5_;K|(95-iesfFc{`=bB z5vku&6MFfU&?~frUa=+gN-d$^+7f!@me8xTgkH5J^lB}kS8oZuMoZ{5TSBkZ5_;{H z(Cf5>UbiLmdM%;XZwb9YOXvkHp%=D&09im z(Gq&gme5$Fzh#wk7m&EuoKZ34KCK=o4E)pVSii@Gh0HR)e`#bmeA+4gg&<=^m#3z&uY)9-H6NJpjePc`L zn_5EO+!FfMc7cdrp0=OEFV+9r62Dk~tR?ivTS9-LCG;m-LVv0y^ru@wf2Jk$XInyl zt|j#6TS9-KCG>4Ap>J;qeMd{^J6l5E)e`#dmeBXKgub^W^nER%?{5kHKuhQcTS7n7 z68hnm(2ul)ezYa@V=bW{Z;4;1|E(o{u708=^cPz~KiLxcsg}@Bw}gJCCG@i`p`U9B z{d`O47g|EU*b@4smiYPl-&?}*|7Z#QrIz^F`afIZXX-Dvg#JoP=&!cKPuG4w^ClC^ z&CTS;t2*}PdA$<8B|b+lY4inB7Kn|9$QEI{Ga&q1}FFw;{bkP3RR{La)>k`mHUYS8fTtN=xWf zTSBka5_pme8BEgxW$RZHltTS9Ns5_;Q~(A%|y-o7RD4lSW~ zYze(nOX!_jLhsTNde@fFyS0RNKQX(n(0kN`-m@k2UM->bZVA0lOX$Tdq4#YGyagg&?>^dT*w4{Zs3SWD=`TS6bv68gxN&_}g|KDs6JF)g8wZ3%r` zOX%ZULZ8qQ`oxydC$)q=xh3=|Eul|s34K~i=+j$5pV1Qf%$Cq+wS+#qCGme3cqgnoNV=y$Y)zPKgyB`u*ZZ3%r@OX$m6LSNAm`pTBj?`#SE zu9nd6ZV7!=OX#axLcgaa^m|)Ezpo|qH7%j9Z3%r{OX%xcLf_C5`o@;fH?@Smxh3?i z?E(>hKW#sUkJSI#5`R~JtR?ivTS9-LCG;m-LVv0y^ru@wf2Jk$XInylt|j#6TS9-K zCG>4Ap>J;qeMd{^J6l5E)e`#dmeBXKgub^W^nER%?{5kHKuhQcTS7n768hnm(2ul) zezYa@V=bW{Z;8LH|E(qdx_+W1^cPz~KiLxcsg}@Bw}gJCCG@i`p`U9B{d`O47g|EU z*b@4smiU|c-&?}*|7Z#QrIz@s`afIZFY7P2g#JoP=&!cKU)00%eRZ)P;*#mHkId!d z#|x*I7fy~XxnEo~?JK6o$LCkQ{JeaPuAh~me4D=gkGg3^r|hPS8EBqdQ0dv z?4;&;%`Ov;uhkNI?UvB%w1i%_CG>hNq1SH-y+KRp1udZ$wuD~P5_-dy&>OXc-nb?7 zCM}^iZ3(?uOX$s8LT}L$ddrs3TeXDVx+U~BEupt<3B6rQ=BU?fr)e`#Xme9wvgg&+<^l>erk8cTmLQCiq zTSA}I68hwp(5JM7KD8zEX)U2oZwY-yOXxFOLZ8(V`s|j_=d^@AwEuqhE34K9J z=nGroy!xV+c%|R#u(+4Mq?aSS{!3dzU)B=(@|Mt7w8Z%Qs+XUqyMAu&dLi_!?E(?w z_Z`3dJbl%!pS!wV2z_@;=zCg1-`f)UzLwDUw}gJ6CG>+Wp&x1q{cua@M_NKZ+7kM) zme7y4#Q6NGm!GHmUWPAsy%72!tu3QLa*HtdYzWg>$ZemuO;;QEulAP z3B8~t^um_Vi&{c&*b;i9me3owgx;hj^mr~lo{Nv?;!8dsPVFM$^W(IZ(5JVAKBFb{ znJuBuY6*RIOXzc2LZ90b`WrbHf2H5+@Q=N{i+j6-xBqH77vHU4?=}6p3BUepTf*N# zbX`k0etk>m8(Ko&*b@4tme4o1gdWeumwe6Kx$2!=FND6UCG_1bp~rLaC0{f5ck(>l z^+M=pT0%eD68gE8(9gGoexW7wi!GsFYKfKWe{TuL|Dz@Jms(={zT+j&S8-pkY}W~) zmum_A=9bX?_qFG%=GcW4Q{V@v3r zT0-yK5_*@G(7U#T-mN9{?k%DBXbHV%OX$5?Lhs!YdY_iii(5kP+Y)-emeBjRgg&4p z^nop*4{8a0a7*YzT0$S%68f-~(1*8#KB6V`ku9N*Y6*RGOXy=-LLb`_`nZu3eP>JPyIMlu-4gnqmeBXMgubsO z^!+WNA7}~vU`yzST0%eE68e#r(2ur+eyk<*<1I0MU3~ny_%d@hg~@&0vRx+J_bt~F z`pqq&`)@oKpIjSr*T=eDCR{)3wS->3CG-X@p%=7-Uf2?PQA_CY>*C|N_;@b94Td>`!FdV5#*c8Gtj-_sJ_ zpZB(eeqT%IYg$5I+YtjnfHhX&?U&^uB z%YEYO9Q))_j?G@@r}ue5~z)t1n|))M;HTSEUv zOX%Nh34KdT=-+Ay{o5^}f2Sq%@3w^gP)q3FYYF}PEusIQCG;P*g#M$J&>wCI{l_h# z|D+}KpSFblvzE|*-V*vRT0;M2OX$C93H{eCq5q~O^xw9G{=1gYA885w_bs9Sp(XS` zwuJttmeBv)68c|SLVvU+wEKc>P53+@92-9u#?OUU+vmdK$*uq8kBh(9=fd!DQ4h}% z^za-(J3k=w2UgH-WTZM zeSsc6F6!aqq8>gj>fz&}9zHJW;p3tnJ}&Cvgj>fz&}9zHJW;p3vd zrR$NdaxTuGZi+cFDsE3b>dic1gKiumj^zd;}4<8rx z@NrQO9~bqX_j(BJeqeTApofo(dic1ghmVVT__(NtkBfTvxTuGZi+cFDsE3b>dic1g zhmVVT__(NtkBfTvxTrlYE->-_sV+M<-WM!6M|ORDx8D~GZwt3{8zB7J$NPm>-~Gbk z$w&X?_Y1!=?|wlK&k^+S96=Ay5%dRoy@Vc~Bk18df*zhD=;1kn9-brU5B7QqJv>Ly z!*c{ZJV(&Oa|As+N6^D_1ULy!*c{ZJV(&Oa|As+N6^D_1U)=Q(8F^C zJv>Ly!*c{ZJV(&Oa|C@$*CU~a_XT=*j-ZF<2zq#qpg+{>A@uMZK@ZOn^za-(56=gHo+IetIf5RZBk18df*zhD=;1kn9-brU;W>gHo+IetIf5RZBk18d zf*zhD=;1kn9-brU;W>gHo+IetIf5RZBWRBc`uN&3z9sYOde-lQe;rY)g2YYDx1OXw|HLT}j; zdaIVuTepPXrX}>YEupt-3B7$w=p9-@@7NN0rp?7Nuy?aaOJz7HV z*%ErMme6~*gm(TgJBQGVYeMha5_-Ru(EGQ9KA=zMv)ag)MPQeNjvJ_XghH68`^p zw1mF6CG;gNp)YL-eOXKB%UeQU(GugoFYxl`QMrC@?s_5gt?dF4-8eP2uH`&&Xk&=UH=me3Eignqas^dl{yA8iT!SWD=~TVnk81z!F= z4L{TM^X0A=LVu+t^jBNrr|a=t+;uQ_U-FigaNn|gOXw9^La*2odZm`oZ*2*^a!cq{ zT0*bd5_+|k(5ttEUZW-Snk}K%Y6-n|OXzi4La*BrdcBs=>$im7pe6Kzme31ZLN96w zyvUe9?mfnY%#Dk2 zAJG!}$d=GYwS+#pCG;^Zp^t3|eOycE<6A_Euk-M3GK0G_W1L~`m&mMtIvynq1V5%*CV|Bt6D-| z-4gnmme7Bj&f^jBI!f3+p_ zeJ!EyZwdWCOXvq%LO;|J`r($)kF6RFu7ayM& zUv}GEp^t9~eL_oU z=kW0?KDjw2x5uJhB-}oKooDg4`nfPoqJc9+&8b+ z%Y^&t^;<%3&=Pt)i;vHXA29dsFnND)4j<3rb2o^Oj}v>5&?mKoKDj0IDJ`K-Z3%r^ zOX$;ELZ8tR`plNlXSIYryCw8FEuqhC34LBm=<{1bU(gcz!j{mBT0*~*&f?C>f1PLX zcly_;xc$eo_;?n7WbVUp^0~sJy-fJr;jxy`kGF(=q9yc`Euo)k3H@|SjL(aYXYpm{ z-W4W~1>^JLEur^p3B6ZK=)GG)@6!@`-`{w-*Vw&w4GP*KY~EK}+ZjTS9Nt5_;p7(3`Y`UeFSH)0WVi zwS?ZhCG-|8p|@-ay;V!-ty@AnhmU9R$;~mjJ;t;6_%$l$%)L8I-XG5GWx~hHc`c#O zZwY-tOXv$*LN96weNjv3ALLp5tNr+YlV|bqYgESHyZPnb&k1Myue5~zYD?(*T0-C7 z68eFb&=0nReyAn%!!4m7X$k#kOX$a1LO!v@(qx0U5Z~xHr+T-)C-821mbHluAv$yBk6Z5X!`X9E0{zomLf2t+)Pq&2r$1S1%NlWN|+7jCBw9v$dr+V47AGO=!pY`h| z{Q7^UCG^j>g#JiN=$~r|{qrrMf1xGxFSdmKrIygY+!Fd%T0;M7OXy!~3H|FWq5pYH z=-+4weNRj1-)sr}TP>k~yCwARw1oa>OX%Ni3H>ixLjTK_(Eq9>^uKNi{d+B;f4?R4 zziA2mZ(BnDyOz*@&=UF&TSEU)OXxpt3H>K_Qgi*My-c|NXDy-sye0I=;0Yb56=ktr+T}D9-a~O@Qk2`X9PVwBk19M zfgau$=;7m{9zHJW;p3tnJ}&Cvgj>fz&}9zHJW;p3tnJ}&CvrT-3wIMLm36)WgR`{n6eIp@)x)dic1ghmVVT__(NtkBj>Edb@-kJ}&Cv zH@NrQO9~brTaZwK+7xnOQQ4b#%_3&{~4<8rx@NrQO z9~brTaZwK+7xnOQQ4b#%wa3K;CO$mX%dU<01xwDzZjT@K`-0*7!acnY5Pt3B{lbsl z{ldb@fBiS!FZ|`a`vpBbBk187K@ZOe`X_t4gdUy|^ze+Jhi3#mJR|7g8A1P4Z?-%s&jG%{S1U)<>=;0Yb56=jCct+5}GlCwT5%lnk zpoeD!Jv<}m;Tb^>&j|XSZbw27?+f(sjG%{S1U)<>=#Tbx2t7O_=;0Yb56=jCct+5} zGlKrT-Y%hsX9PVwBk187K@ZOedU!_A!!v>&o)Pr$jG%{S1U)<>=;0Yb56=jCct+5} zGlCwT5%lnkpoeD!Jv<}m;Tb^>&j{M%!a*iJJk`st4bKRAyf0YtJeu3%hwZ*#{McV| z4)EXiN8w|CDIe@{&*a~8|MF{hzu>+OzR&FYxwdpo=y$b*es@dg_q2pwrX}>UEuoid z3B7zv=oMN*_s_&xuhnpW{Ub!XoDlMT`Z3(?vOX$^GLa)&hdd-&5Yqf-4yCw8G zEuq(K3B6uR==EDdZ_pBY!K-Yze(pOX#gzLT}R&dfS%J+qHz=z9sYyEunX83B6NG=$%_a@6r-_*Ot(`wS?Zi zCG;LGq4#VFy;n=(-L~$meBjPgxJpjePc`L54D88sU`HyEun8|34Loz=-XOC-`*1Xj+W4OwuHW`CG_1bq3>-M zi1>?X`!;;6{;QVo{L#l-LVuzq^e0)Ar6u%NTSDL068ip@&=0hPey}CiLO<3L`tg>~Pqc)7 zvL*CWEuo)oi9f5q))M|7>N72&pKS^KTubQZTSC9k68goK&@Z)wez_&|D=nd4Z3+Ea zOX$~I;!o?Zw}k88XbJtzmiXiPTP^WN^|xC>f2Sq%cU$5QYiH+=nJJ5>yXoZGlKZ~R zrhV=7`X)1)n^}4-G zxV~OX==EDdZ_pBY!3B5y0=p9=^@6-}{=a$gBw1nQZCG>7Bp?7Zyy+=#vJzGNW)e?H| zmeBjOgx6qgz5B(-QjFme9wwgg(9{^a(AYPizT&QcLKQTSA}G68hAZ(5JP8KD{OM87-mD zYzciik8q(?J>3E%&NEupV(i3{r+ zT4KHW#+F!j(EeP|_thWj+X-*)rk2n*w}if>CG@Q=p>JymeS1sjJ6b~D*%JD$me6;% zgub_3AY$!l`!=jq|5Zz@S%17G^e0+Ef3hX?r&>aPx+U~yT0(!eCG_W7LVvy`^cPw} zf3YR>ms&!9xh3>hT0(!dCG>qQq3>@A{Xk3T2U|iv)DrsPme7y1gnqOo^kXfdA8!f$ zL`&!=TS7n868hCG=V?q1SE+y^fvKT(8^9gzM|IgkHZT^ad@V zH*5*LQA_BJTS9Nr5_&;P=uKNfZ`Kle^On$Cw1nQWCG=J;p|@@cy-iE#ZCgTb*AjaB zme4!2gx;|w^iD0IcWw#2OH1foTSD*F5_bZVA0lOXz)DLhsiS zdjFQt2egDfuqE_CEuj~-gg&?>^dT*w4{Zs3SWD=`TS6bv68gxN&_}g|KDs6JF)g8w zZ3%r`OX%ZULZ8qQ`oxydC$)q=xh3=|Eul|s34K~i=+j$5pV1Qf%$Cq+wS+#qCGme7k@LSNJp`r?+*m$Zbwv?cW7mUz2oacAW#dpm@;e^pE9 zt6M@}(-QjHmUz2oac}R2-VQOI#mBSw{kQ!SyNZi(?MKAy#woqJc9+;=Y5%Y^&BLT}s>dXtvW3tB>N+7f!Rme8BGgx;bh^mrB@&*I}*e96bdcozTfKQF#- z_ZL@mQxWbzew5DQ`}FI*rC&GU*MDnE>|Nj160YCg68eso(08_kzN;nl-7TT-Z5N0b zpBEpW7k{Sr1H$vj&$fhqt|j#IEumj%3H@SA=$BeTzuXe~m6p)2wuFAICG_hp@q>C^ ze5tt~{(GOZ%NP2&dCq)xO|yNj}5;6 z?Aq^4t}Xe0uD?6&+3oXN)1Lj@e{$HSkw(0IZCXNa+Y)-ameAX`gx;Yg^o}i|cWMc}b4%!5T0-yI5_-3m z(7U&U-lHY-o-LvGY6-n}OXz)CLhsuWdcT&?`?rKXpe6KyEujx;3B9l-^uaBm4`~T~ zXiMnBT0$S*68eai&_}j}KB^`3(Ji5mX$gI7OX%ZTLLc7}`h=FyC$@w>sU`HuEul|o z34Lly=+jz4pWYJsjF!-6wuC;bCG^=Xq0eaveQrzW^IAfm-xB(Qme3crgkIDV`l6Q5 z7q^7Iq$TvFEuj~;#69(8E%6)m`&+{Izq}>%6)mB!Yzci;OX#axLSNGo`r4Kl|NgNz zes1ITxx3qi(D$|rMEvr!eH(tM{;QVw#roqdp+C_Q`jai8Kh+ZY(=DMt(-Qi#EulZx z68iHkp})`)`im{0ztj@?%Ppb5(h~ZsEursg34MP{=m%OtKiCrbp_b4Ow}gJACG?{$ zp&x4r{di00Ct5;3*%JDxme5bP#4preYl)w$pJ@sGY)k0pT0%eH68eRf&@Z-xeyJt& z%PpZ_X$k#mOX$~HLciV;KVN^nC0zeTOXzR5#7F9HwZzZX-);&0otDtwZHb?$|9fBa z^LEeTzuntjs<%$~SY5g$^t)O@zq=*$ds;#-(-L~wme9+!gkHWS^a?GZ`)53hPi~F5 z+hfCCCfq(7wS?ZdCG;jOp%=7--n1q3W-XyNZwb9cOXw|ILT}X)dh3?Z+q8t3B5y0=p9=^@6-}{=a$gBw1nQZCG>7Bp?7Zyy+=#vJzGNW)e?H|meBjOgx6qgz5B z(-QjFme9wwgg(9{^a(AYPizT&QcLKQTSA}G68hAZ(5JP8KD{OM87-mDYzciEukN63H?}0 z=*L?^KhYBU$(GPhwS<1UCC0P(cotuF?pC-|F8#_T7Fx<5_%s zUi`kf50A-Nc07xZ&x?O|`u;FJFFu~dH<-RJcBqaZ_jAIJe_KoF+gn24(GvR3me6;#guc5a^u6r@5#!gWj9;VjOz#JTujzQU zCG>MGp`ULF{X$FV7h6KV)DrsTme8-XgnqRp^lL4lUvG)=YgERsQCV~DU14(HwpK3_ z?(5cW3B68B=yh8{uh$ZK{g%+q;p16+a&t^>j~#lEaQlp3qcWbw$FEUY^6_wCFA_dJ z7PW--oH(H`ZV7!!OXy2mLN9KKw|f@Xy*gw1mF4CEjUYFaE#w z^Sim97vbl7OH1hSEIxj{_@i_04wL8MAM0hp*O)xs68edj&`-96eySz((=9Q6jmr2n zD$CBjD@^Wtmg{B0ebMqQp;u@L-9Hm&y<#sAuCLS*dgYeTtF(k(wI%dwEumL$3B5*3 zXy@?pEIzq8Cb!3yy-2uywrUBzbxY`NT0(Ez5_-Fq(A&3!-k~M*jxC{gY6-n_OXyu% zLhsrVdbgI)ySIeiqb2m7Eur^n3B7kq=zUs3@7of3zn0MZw}d{RCG>$Up$}>ay|5*; z=fnwpNK5ELTS6b!68i9#&_}d{KC&hBQ7xg5ZV7!%OXy=;LLb)>`uLX6C$xkq0ep!eNId0b6Y~6*An{tme3cpgubvP z^rDu~f1PLX|I*vLthY+ceaGSt0nZ^Eursi7l;_27ayM&f2Q{X!t>70wuFAJCG_(x zpdXtvW3tB>N+7f!Rme8BGgx;bh^p-85w`vK! zbxY`NT0(Ez5_-Fq(A&3!-k~M*jxC|bv-o%xAJ5`TJ{}hJBH`oXqL$DXw}if=CG@2& zp%=Hr+kIZ#+q<&2L+GnoLSNky`kI!|!?XC8>g)P;;;nwYxVLvhuSdMq-@Ez0_4+sU zdW5$(p2f#MN9EzUcZbRI)Q|Ks;d$#vTS7n968iC$&`-34ezGO>Q!SyNZi(@E@$q@_ zW#`@%CihLt^)lhUYWbGXE3|~}pYeI|$*nPWdyHrC@hpD8;QisiUM76J9MlqeVN2+P zTS6bw68g}V(1*2zKD;IL5iOyQYzci-OX#CpLLbu-`q-Ax$F+n$z9sYtEuo#m$Fun4 z=9t_bi+Yi8`}{DU#ouYqi~sN4-s4$(Jd2NK@wxjS&uhQji-dlqCG@K;p+csR2c2_GM4wS+#qCGme7k@LjTV^i~k?}c$fF%5`O$ET0&ph68fr^&{wyFzNRJg@Oklf+F9JsZ#;{S zXYuhYKKJqJEd9f97GG-amoxeIEz_6(&wg$G?D3_iZ7!L<{9h;6mi+&+|7-evhi}LI z82;}!c{{WF%m3E3_sxB`AN$Bs3-$ZemuO;;QEulAP2|fJ(4|=1%ozNS%gx;hj z^n#XHtKPIFT;HrE^yV$0w`d8yWlQL-T0(E#5_+4K(A&0z-mWF|_AQ}zXbHVzOX!_i zLhsxXdY6{ayS9YhttIsCEur^l3B6}a=)GD(@7)r5pO(=3wuIiVCG`F+p$}*YePB!I zgIYo_YzcjEOXx#dLLb@^`mmPJhqr`2q9yc^EuoKU34L@+=wn(!AKMc8xR%hzw}d{S zCG?3cp-*ZFeR50aQ(8iw+7kM-me8lSgg&Ds^qDQ8&uR&Mc1!4UT0)=O658iGgg(C| z^aU-UFKh|Ds3r78Euk-N34KXR=u2BdFK!8aSxe~mw}igDCG-_7p|5NSeN{{7t6M@} z(-QjHmeAMDy*T;z^3$Gu+nV(UxJvleRE6bTUtWj z+7kMQ!SyNZi!XuueHQV^)oG@pKS^KTubQZ zTSC9k68goK&@Z)wez_&|D=nd4Z3+EaOX$~IV&(elE#dk%T0(!bC04Az)e`Tmzugl0 zJ1wEV+Y&3(!*lgwIc4p6*Y=%VXxEu{ZNcE$y3=dR4!&)6ef_R2Klt|N-@D{&?3?vU z^R6$L*R%Tu-~5H%*z99nuQ%`Kb%k%7vcbG-H_k4!8_v7t(&f|Z8_l~me4F$CUfJ2} zc-4*Py`FEIUEgHhwc+i)VBWRMyyBEi=UuyYcA zEuq(G3B7Jh==EAcuip}SgO<=6wuIiOCG^HEp*Lv>y`UxZrY)g2YYDx1OXw|HLT}j; zdaIVuTepPXrX}>YEupt-3B7$w=p9-@@7NN0rp?7Nuy?aaOJz7HV z*%ErMme6~*gx;qmwEOtk{k-0=zMv)ag)N~MwS>N?CG^EDp)Y9( zeQ8VR#Vv79eOXJq)!$R??OoaH5&Ei@&{wyFzNRJgwJq_3{(T&NJ~#FAAoR^Gp>Jsk zeQQhT+gd{3-V*wbme6;$gubgK^xZ9??`;=|7=KUo8$WmR`yTG=b|LirEukN13H@M8 z=!aTDKim@fk(SVpwuFAHCG_Jhp`U08{bWn%r&>Zk-4f&P<9OreZvMX8Z+E*8`a3P5 zzuOYypZmY$|D(7sS+U!M&?~isUb!XoDlMT`Z3(?vOX$^GLa)&hdd-&5Yqf-4yCw8G zEuq(K3B6uR==EDdZ_pBY!|P{%e4Nt~`rMY#=e2}Bza{hqEuk-L3B9N#^hGV9e~@SKxB6#| zex;x9_dAP^UpqFQ#lO`1HQ_A!<(AN2X$k$+meBXLgucHe^mrCu^8ZnME!Yd)E`)xu zCG<-z@x$`9V;`IQ^;-BwHW776#wOSgo6S4-%3w}gIAOXy`rKX&Y8!u!WgEunXA3B5~8=v`Yv@75A} z_mZgg&q(^g%757q)~xxFz% z34K^g=)+q=AJG!}$d=GYwS*p@7yscoi*M1L@n!vYC4SJ)i;rjV@hm=`#mDExzcF`r zG&%c!vzLjF^!W3wmKdKGAJ5|B^Wt+K3(jWS_9CITYYDx5OK9it@hm>MIVQKqcorYO zM&)no^Wr|9uk2l!&{wsDzPcs!59)dG>sJ5qeLvcde`7xnV#^iZ`1=d|{BG*o34L=* z=v!Js-`W!TwwBPhw}if!5eRoUfd)oyf#;;KspBI0o_XEQ7*3Y(tey%0- z^DUuYXbJsdOX!zcLciP+`jwW@ueOAKttIs9EiwN2M&tA1YtFqZOzzu!eG1{eZta%P z>$HS+4j<3rlbd65du-c_gxhDkmeAX`gmw-m^!U8^_`LWDbMFq5_lFaEnehH`QcLKQ zTSA}G68hAZ(5JP8KD{OM87-mDYzcifIexwPrrn1T9ve1m3B7qs=q*}8Z`l%htCrAPw}jrNCG>a}AHPQBn7Ma{$@{~x zy-axjIIbo1@hzd9!^gAu&5-}*Yx8O`r4KlzeZ(zUVJ=@k7x0}jc4)k>&3^j_;?l{zh3;1xueYFY;$NY6V5t^ zwS+#rCG-(3p^t0{eN;>6qgz5dhmU9R$;~mjJ^nU5FaCDV;{N+z*)N^YSG9z`x+V1A zOde z-lQe;f|k(ZS$sT;k7w~E9}nZ#sQiyVFaBS;|G2E1ju?OM=6DwWRPP6bv+;NqAOB3l zm*(CbCSULUaxW9(*NZRty(i=I;^Xt;<5_$>i;rjVxyJ^dn~rDk|M9clNSl zcSE_ouk5BL^i?gPJtt1+zscvtz5Vfd@g;vZl(Y25yWI)>iI&ixYzaM{#m7Hi{KdI< zhsjynIsAw1EWXtA*JozGk9(W$+lGH{l)BEpV%s;lws`K+k8SrR*Jf{X`!~6E{q*&A zc#~^aPp|FxCf6>XUfbzSuFZbDo!{iz?8n>XO|H%U|KYB0a&7iDcbj+3+nc@3-RE7q zboz5Nd(68w`?^cl_MBc@TEv3Mzvuqt?>l|pp!b@7-&LpI_TiTBcZdC*me7BY6<B zv6j$((GvPETSEW)meBv9CG>x63H_g1LjUKM(C!O%HsRxdaBcip7(W(%)IJs#PTu?9 z__*lrR~fsqd56=kN`GC+r*%EqqM$p4Ef*zg`^ze+Jhi3#mJR|6z?#ChY@Qk2` zX9PVwBk19Mfgau$=;7m{9zHJW;p3tnJ}&Cvgj>fz&}9zHJW;p3tn zJ}&CvABWJx$3;DST-3wIMg0f89--Y2%H@NrQO9~brTaZwK+7xnOQQ4b#%_3&{~ z4<8rx@NrQO9~brTaZwK+7xnOQQ4b#%wa3K;CO$mX%dU<01xwDzZjT@K`-0*7!acnY z5Pt3B{lbsl{ldb@fBiS!FRU@|enAh<2zq!%(8Du={>k1hp@(M#Jv<}m;Tb^>&j@;W zM$p4Eg8u1#96}Gz2zq!%(8Du=9-a~O@Qk2`_X~P>M$p4Ef*zg`^ze+Jhi3#mJR|7g z89@)v2zq!%(8Du=9-a~O@Qk2`X9Rsuw;Q2{_XT=*M$p4Ef*zg`^hbL;gdUy|^ze+J zhi3#mJR|7g89@)v2>SQ?aR@y;Bk187K@ZOedU!_A!!v>&o)Pr$jG%{S1U)<>=;0Yb z56=jCct+5}GlCwT5%lnkpoeD!Jv<}m;Tb^>&j{M%!a*iJJk`st4bKRAyf0YtJeu3% zhwZ*#{McV|4)EXiN8w|CDIe@{&*a~8|MF{hzu>+OzR&FYxwdpo=y$b*es@dg_q2pw zrX}>UEuoid3B7zv=oMN*zqciH|9tQ7i`6UDgkHHN^eQc(S8WNsT1)8FTSBj4CpFh= z_A=r6S}md1ZVA0kOXzi5La)~ndi|Er8?=PpuqE_HEulAV3B5^6=mjm||3hus60W~* z?#0Q!_pUeVCBpU1TS9Nq5_-#)&|9^H-nu39HZ7sIZ3(?yOX%%eLhsNLddHT~JGF$~ zxh3>2EunX93B6lO=-pdF@6i%^&z8`8wS?ZgCGv(66?Heyt_+>n*Wz{q>e`{TnTzzu6Ki*57K0_txKT3H_ax(BEx|6>4Ya zKQyt_R4@D3TDm3lyIMlOyCw8{T0$?=5_;K|(95-iUcM#t3N4}E+Y-8e#(0IZCXNa+Y)-ameAX`gx;Yg z^o}i|cWMc}b4%!5T0-yI5_-3m(7U&U-lHY-o-LvGY6-n}OXz)CLhsuWdcT&?`?rKX zpe6KyEujx;3B9l-^uaBm4`~T~XiMnBT0$S*68eai&_}j}KB^`3(Ji5mX$gI7OX%ZT zLLc7}`h=FyC$@w>sU`HuEul|o34Lly=+jz4pWYJsjF!-6wuC;bCG^=Xq0eaveQrzW z^IAfm-xB(Qme3crgkIDV`l6Q57q^7Iq$TvFEuj~;gubjL^!r;vU)~b>ik8qG)hAp8tY6-n@OXy8nLN92E@p*?D$$_Wm$C zH+%nBe%iBdTeBX|(n}r>#~n4}@7sNg=b&fT7RV|4&Aax# z*@brhdDre4Tzme#V#(W_ea`P&zq-O3ugyL`_pJxadz-V*F+4B>=%lqcNo^PA|nE!41 zdLQ$RQ)X}XcP7`CJlFHP)1Ezk{MNK*x5;l#d-iMn?b?6e*?(hy=EKr8q2JXK`rR#| z-_sI$nU>JYwuD};CG_$wp}iYS-W8_ZH+}E_Ex7Kcn%rE|>npW{>npc}UZo}Usx6^? z?lk-SNe@4N(!!_S}e@bf1<{QOA|KY!B0&!6=0^Cvz0{7DZ# zf6~LxpY(!$y@cMhCG=)3p*L>{y+uptEn7lw)e?H^meAX@gx)7eOF89yIVrv+b$6C%hUF4_@(-JcwuFAJCG_(x zprme8xUgkHTR^cpRp*K7&B zR!iu$TSBkX5_;X1(Cf8?UcV*u1}&jCYze(lOX!VTLT}O%dO=I*OLpU@Ke#Fo$}wS+#oCG;sRp-*iI zeOgQC(_2EH(GvR1me6Olgg(0^^f@h|&us~PUQ6inTS8yZ68gfH(2H6^U(^!%;+D{t zw1mF2CG_H!xTn6XCH}w8;)kvH#;^PF_O9&p2z^yc=&M^oU(*u$+Lm~y{~fuX>*sq* zKTpEX|JIi9_k!Kl60YCg68eso(08_kzN;nl-7TT-Z5N0b&*I}*{F&Yl2xr%4TS7nA z68ia;&@Z%vez7I=OD&;aZVCNLOXycyLci7$`t_C=&*J0r;%m;mD@^X&*6L-#ecjqE zq1R~%y>3hB^;$x&-x7L*me3ovgx;tn^u{frH)#pIpe6LCEulAS3B7qs=q*}8Z`l%h ztCrAPw}jrNCG@r}p|@)Zy?smQ9a=)~*b;g?i;rjV@hraN<6%)R5Of$Q_>- zzrXhb!t=}zw1j@JCG`uLX6C$xkfpBI1U zJ}>US&lUZ5BJ`Clp|5HQeRWIdYg$5I+Y*11&x>ET`j7AX(SE!e`*DaPR(#{{kM;Ap zsc$Fr%`KsCX$gI6OX%BLLf_sJ`i_>+ceaGSt0nZ^Eursi7l;_27ayM&f2Q{X!t>70 zwuFAJCG_(xpBo+HlWVj4tmEJ0+U&OTQm%>6*~*Y6<=BmeB8M3B629=w(|%FV_-!`IgXb#>q`NZT~qY z|4F9ZclqAe=llLE)r4NTCG;vSp;v7Qy;@7?)muWZVJ9`$YxXkX`tbil=(YQHLa);j zdfk@L>$QYlza{hrEun{>`{<4Oc0zC55_*%C&u{2C#*&#v7rgx;+s^zJR8 z_h<>dXG`e4T0-yL5_+GO(EGN8-mfL}{w<*oXbF8_OX!1ILN9CyeQ--`EoRLoK0iY6*RFOXyo#Lf_gF`nHzPx3`48qb2m6Eursf34M1<=zH4* zBF5j({Kn6H{5g~Nb-NJy{+7@Ww1j@JCGOad z-l!$?#x0>YX$ifcCG@5(p*L#@y?IOMEm}fv*%ErIme5XE%A2$|IBxLdpGoUi1BNs#;=vSzxM;e*B(9468gcG&=0kQez+y{BQ2pHZ3+EY zOX$a2LO;$ZemuO;;QEulAP3B6%U=#5%JZ`=}kla|m6T0(Ey5_+?i(3`h}-l8S+ zcorYe;^SF-$;ZQ)y-4`@IIAV}*)5^ZX$gI9OX%}jLZ9Cf`hu3w7q*06)Drrlme4=Q zv$)?^@J>35yFJFU`1t4f+&A~(F*(bQXYuiCrM^3TfB2jH^L)m$_;?l{&*J0r;^Xt; zb9c7RVyE^Zp-*cGeR@miGg?BAXYnOptK;^#yxWENe0@brxc;Mb79amSpYhMue5Ch# z!t>UTw)}tCy9+4E@5A2m7E{E`ZON8d%s66ZX66wyE8w8R%oA|XNgT^zdQSl}h#4Y- z%yw)q@1C>w)7_={O|Lg+pX0ZFvFaY(`PS|K^o*ybyXx0HJ!9x+T0%eD68gE8(9gGo zexW7wi!Cwzy!f}~^Ww|R-JP%e`I$WTyL`L!|2IHQ{yg#PtN%Z*-Rwni-h$Vg{28BX z=P$T6c^%g-Sa8kDCf6_g2G=H!!595I*DhXg?J_^;JeMrE=EajgAH8(JwHqg|)L!=Q zT)TY1weiQiV*c6#^FNmRd~(108pP#mLa)#gdc~H|E474Pxh3=}EumL!3B6iNXuqYi zzpeA#yV~^c7<3=N{j)_&=q+19Z`Bfd>z2?nEpb)-+LmyAo4K2_-(Fd7+bzO%e||t{ ze||vd?OH;=p(XSiTSC97CG?wHLcgUY^jlj(zpW+o_AQ~uujkR@*YoJ{>v{Cfy&j== zX$ie+OX%HNLhs%ZdXJXSd$xq$t0nZ_Eur^m3B7Mi=>1wk@81&ofR@mUT0$S#68fN) z&TEuoKV2|a#2j~>6CM~`36 zqfhGf2<>AQLZ8wSdi=OWj~}<_@#7Xf{uzQkvmb}hXSIYryCw8FEuqhC34LBm=<{1b zU(gcz!j{k%wS>O7CG;gNp)YL-eOXKB%UeQU(GvQ~me5zVguc2Z^x~G#*R+Jbwk7m+ zEupV(34KFL=o?!?-_#QN=9bX6w1mF3CG>4Ap>J;qeMd{^J6l5E)e`#dmeBXKgub^W z^nER%A8cntO#hC-^zRtF)Y}2!|KIl$EulZz68cjup+DUc`sJ3;pJ@sG*_P0sYYF}N zme5~l3H`;E_;&sO53fA;)4y_VmHv*w>hmw3^qLELWuKI!oQa=d7N_Xx&^N{c}#ch`UTg?SoeCL8|UN(9Dz3UrXoADju?)hu)oy*ye>6V+n z?QwtkZWDTime4Dv|D{n3`tA8QHy@s`jJwS<1SCG;aLp&xAt{a8!r$6G=_(GvQ}me5bNgnqgu z^fN7?pKS^KTubQZTSC9k68goKnEqYzW&b?tKl*pcm(1R)bKfldyX2GezN`25`!IfN zGG6=c1=l8@Gx8eWv*6mLUO3P9F1WUMa-rRJ!L`ZbmG83fW3S2QoLu|91+O>x{FZBP zUvO>wHGKbqYva%N0}HNA9{c^6Ke*u9V6^g$t|HzX$gI5OX$;DLZ99e`iz#)XSRest0nZ=Euqh834Lx$=<`}a zpWhPtf|k%1wuHW@CG^EDp)Y9(?Kx@ky!2!BWi{dY${jV=8g+Vs!HrXMrC)cYUd&z?Tf68e)Zp+D6U`qM3;Uv3HgnU>I>Z3+Fk zme8MX3H^nZ&|hqcZ`YrVJuvsv&z_s7XXm*)!(;1u-6XVgGohzv=jqp#?K=17F#Gnf zTek^sAG^1N-lHY-o-LvGY6-n}OXz)CLhsuW+PQgpcAmXEX77*Vx=FZyj&BKlLQCiq zTS8CI&fomm+5Pt)IXnM&uRlFIPrt5g`sWREk4OHz;Mr~x`ni_S&$onrp(XT-EiwJ` zhHuTEH%!mYbI+CQ_BKU$?4O>U=k5#N?%DmZZA-X+Ue^-(^(~>dYYF{^me6l(3H_#) z&~I)D{g#%{Z*2+vwwBP_w}jrICG?Ihp?7Kty>m!5eRoUfds;%@+Yo%dk z-V&?Vs%&}O@7RuTkv|5$8Fbse!;cz*ZvC&u5IHL=J~}1*CvmxF8|VkYhE___5boWxHkDU z{>uEd2fRY~-2e3R|Fd_+>|^|T-6TBbuWx5-t~cm5;rfOxp*Ly?y>UzEO4+7f!R zme8BGgx;bh^p-85w`vK!bxY`(mYDwigXzaG2ljSAc>6f0CG^29p$};Z?c*0hPd|qF z&i&e%U+wqzl762Fzu%r$CeJPUf8>v0{O_B7472p{fBG@Z^kbOm$1wlxAHz(~&U26b z>-N}Bcx;`Xo#*Zg-%daOKmGjw^z1zSy7PQ!SyNZVCNNOXz1?LO<6M`uUd7FSLYy zu_dNoJ2U;-ncgI3@mRQKw+OvfOX#&*La);jdU|&LfB*MPr)THs*?D?)o}Qimv!0zj zzW@JucK-E#zFT`OCCG;IFq3>)7eOF89yIVrv(-QjLmeBXLgnqD{5i$L| z^YrXI{k-#k`{$jPn|t?p<B`<-^A48#|C=**g%gT8|d+413i9hpx5ii zC$z6eA@uk)DSG^x6unX3kI);pgx;hj^rkJLH){#Kc}wUmT0(Ex5_+qa&|9~Jo@ojF z+LqAUwuF9NOX%0Pgx;lT0+0ICG^``LT}#^dWV+KJGO-0 zsU`HzEunX53B7Ae=-paE@7@x6kCxDTwuIiRCG_4cq4#MCy>CnC{aQlr-xB(Ome7k@ zLLb-?`k=qp-6U)d7+s+Q1Kw}f8Y68f5!(ATzv zzOE(o^(~=qXbF8|OX!;sjEFy)Z!g2U>VIsBKdk?$CGw3F{qdI254D7TxFz%>EukN63H?}0=*L?^KhYBU$(GPhwS<1UCG;~b zp`UFD{aj1v=UYO*&=UH^miUAEUt8k$>X%wVf1)MyCtE^)swMQNTSC9w68bYOp+DOa z`g1LzKi?Ai3oW6)*b={Af2k#0|8h&{ue8MP*8kQLzf*s;CG^)?LVvv_-dT@l`n}~m zf3o137hf`e?cLwt+QswN{&fD@1M|IHOZeD*`IgWtw1i%@b0MlGQ? zZVA0fOXy8oLT}a*dh?dhTeO7UvL*CZEupt=2|d#i)2|6%_B`(X*}40L(7UvR-nAw4 zZY`mAZwb9eOXxjYLhscQ+WYwA{ao);6MEm4(EGK7-oGXE0WF~ywS+#fCG0))M;gme5DEgg&w*^ieILk8TNlOiSovTS6b#68iX-&?mHnKCvbANiCsI zZV7!#OXyQuLZ8+W`t+92XS9SqvnBLdEuqhD34KmW=yO{_pVt!l{Fcxcw1mE}CGJvleRE6bTUtWj+7kMp&jqej)UuEukN43H^9W=qFl2KiLxcsg}@Bw}gJCCG@i`p`U9B{d`O4 z7g|EU*b>vPYgzU@zDD=YSG!*b{k4|RUvG)Y-)ntfE@#ip%gx{RzO{U}3B5u~=oMQ+ zuhbHH<(AN^w1i%@CG=`7p;vDSy+%vu{+phiXYY-<`(w*)6Yig_T0(E#5_+a3rf28r z+4;cU4hU}_2epJgxFz%34K^g=)+q=AJG!}$d=GYwS+#pCG;^Zp{Hl(>Dl@J z>)F}c%YV*cEup{C65mz-TT4vO&fQ0Icg5^;=~~?;JXWsV5_+AM(CfB@Uauwe`gXSF zdV_8gu5Z{9dZU)m8@Ghsq$TvGEulAS3B7qs=q*}8PtVRvzYfIxv3>Uop?7Ery<dXG`e4T0-yL5_+GO(EGN8-mfL}{w<*oXbHWj zCG_;{JpH`$g>!EXvu_UbxY{Q zEupVz34Lu#=<8ZSU*8h?hL+Gbw#0kun_9y4n_EKP(h~aCme9AggucBc^c^js?`#Qu zS4-%-TSDK{68hej(D${3ez2VpG5x&r^z8gnZwG|4<|kT0f3hX?^z6L!>%O;{yIW?z z{j=_y*Yq{vs|+ujJkH$P3s)On?(4u;U(k~?tQWgJv*7yVG180g>c{hbJmJshf9~bx zUod$*mVdF}+TzJ8wYM+0HhEn0T_#^2zRvvjzrQ7Xea&~bg#MnE(BIn<`ukc!zr7{& z_qT-pftJue*b@4OT0;MDOXweI3H^?i&_CJ|`o~&A|9DGizqgA_+&{0|uD#msi=XJ% zP5AZyWJ~CuY6<<*Eunv=CG^j>g#Njf&_CZ2`WIS4|6)t%Uup^c%PpaQr6u&QwuJt* zme9Z65_(BX=-+4w{hKYJf2$?*Z?}YgXG`ecX$k$iEunwU&emN2ezytN|DYxGAGU;k zS4-$WY6<$UsyEz*1zoY;ve?DFg`Ep z@roJ}>I=c~Otgi+X%s)Z_D_9-kNW_`Imc=S4j}FY57mQ7`F! zBlP&ZsK@6;Jw7k$@p)0dvmb}hhXC|zpI~z(Bt!>9-kNW z_`Im!-RlwBhXC|kI##Gd|uS!^P(P~ z7xnnOsK@6;?Rjy=#QpQS?b`IQVCfmz{qddtSTMdVEa`24@N1twF1-347Z%O_?l1ee z;P0i4j|+M{Bk1vrpvN6(Bm0Fk7oouo)PqTM$qFKL62tyJ)RNtct+6U89|R{1U;S+^ms3$>h_*kIFGlCw^2zopt=y&$x5PCc#=<$r8$1{Q+&j@-vBk1vrpx@QcL+J60 zpvNgW6A8dbKguJ z7be%X8Ls{C>SF1ixy%gL{&N2DWb%40-86swufD;x>*lZh_53xDt?>GjAH$!KE?*OR zg_h7OwuD}(CG^TIp;u`My=qJ7)mlRPCC~n{=X;HoaJ~Pge}7>1-k7^Tw(K_H{@JP} z^wuq*XIf(V_Xnn519f0;2ZXnegIYo#+!Febme7Z`gg&e#^x-X`k7x;fWJ~CyT0$S) z68e~y(8sofKCUJ7@hzcGXbF8{OX!nYLZ93c`jnQ?r?!MXttIs7Euqh734LZu=(Acv zpWPDroR-k%wuC;fCG`0%p)Y6&ePK)Ji&{co+!Fedme7~BgubjL^yMv~uV@KJ*peM?K|TU$ck))M;m zme6;!gub&S^j$5X?`{cwPfO@~TSDL068gb*M#S`Mpr&5~^-^yK#Pn;RrhoT#t+_Xa z*~f~tyG?lPSf?fQx-FsCYYDx+ovpdvpxcD&8@7brs3r8qEulAQ3B74c=*?O}Z{8Am zi||LVu(s^haAlf2<|+$6G=_)DrsP zme7y1gnqOo^kXfdA8!f$L`&!=TS7n868h(B3D>{e68bAG@vZuIYgV57;a>TA3F+_Nu0H?rNl(tD zlb%^{{mP;LW|d|CKJDN2=72?*BS}?Exme4*PA@s#9p?xlY z^7(xKUe5o>Ukl@Z@3-P(lAC6K^>e>|&+(Jzc|ZR2W0L8|B-4*c=H5Yk4DfO{3H_b; zG0F7oJoi|*ZokKb$HMhmLa%RUYpyrwHsSh)EulAR3B7Sk=uKKeZ`u-ivzE}Cw}jrJ zCG_;{JpH#c_nUijn05NRmgeNzjl;DEdpj9Fm%U`T=JVjL zO+H_J({Sw{yLRW$|Fq!8acS~-jLDBP`T1`dUhkh5yxw^2Ulv^Rcsb#H3$A&rpRn-n zxK1APU4Q?A*PA>ZyY_(v*T&cT;DT%8&-bAP*Cx*ye#{RqxVE&(0Jj3-F7v|C1OrrgG{ov3atH-y2@qgn(3$860UT^Zh|AG0J z`Td>z9zWabdApgs?dTP1Li=+LLa)>k+OK``>({H)g!bnggmxcI?kDX&n%qx%{Id?- zf3N&mhjt%M?nk{=P3W~-La);jdi;1z`&f<8>$ij+KVH)t_HsgxAFt{0<25~gyr##G z*R+q-2t9thrpJ%h^!V|b-m2Fl^wuq*XIkR<`n4_L`nD~hU)K`)^(~>dYYF{^me6l( z3H_#)&~I)D{g#%{Z*2+vwwBP_w}jrICG?Ihp?7Kty>mClsADc}+KGWmJXZnDy6M9ig=mT3qAJh{1;Fi#bw1hsi zCG=q}p$~5feMC#>BU?fr)e`#Xme9wvgg&+<^l>erk8cTmLQCiqTSA}I68hwp(5JM7 zKD8zEX)U2oZwY-yOXxFOLZ8(V`s|j_=d^@AwEuqhE34K9J=nGpyU(^!%;+D{t zw1mF2CG=%2p)YRJ#neN#*5n_EKP(h~aCme9AggucBc^c^js?`#QuS4-%-TSDK{68hej(D${3 zez2Vp@yvXC8J@2Hu_d0W|EVSPKevSbmzL1)YYF}Sme3z)3H`yA&>w0E{o$6-A885w z(U#C3YYF}Fme3Eignqas^dl{yA8iT!SWD=~TS7n468g!O&`;Ufn(L>#O}PF{OXz1? zLO<6M`uUd7FSLYyu_d0Y|FtC^uU~2j{fU;)pKJ;Jsg}^6ZVCNzOX$zEg#K(x=+Cu; z{(MX5FSLaIVoN+xf2k#0|8h&{ue8Kt^}n^mqxDx?LVv9#^w(SBk=mJiWfQB-&Fr_% z+^a9>H5c^Eg5GB67kmB5`=J-VF#mHcH~)QR=GQp4`QNdQiTmeu+pqb%TS9+NOX%-y z3H^O7q2JyT`ukf#|3FLVA8ZNzLoK0yxFz(Dw1j>~OXweM3H@U&p?|z3wEJR_iTmeu z+qGBQeeo0hx(UDjpKJ;JQ!Sx?x+V0_w1ocIme4=f68h&`LjOWb=wECJ{Yx#Of4L>} zue5~z)t1n|))M;HTS6~s3H=)_p?|X_^l!C<{_U2~?`#SEJ1wDqwjnRG~xY#aBcd&FnwQmwY@JantkhE_Ic5{ zaC~0W;~7DZX9VqhKhXC|kI##Gd|uS!^P(P~7xnnOsK@6;Jw7k$@p(~?&x?9|Uex3B zqF&PdM(FW*QIF4ydVF5goJ}>I=c~Otgi+X%s)Z_D_epf#ap~vS% zJw7k$@p)0dyVoPM$AQUXfgYb1_4vG~$LB>oJ}>I=c~Otgi+X%s)Z_D_9-kNW_`Imc z=S4j}FY57mQIF4y+VkR!iTmeu+qLOq!O}Cb`{O(Pv0!{#Skl`7;nzNWTzK_8E-ae; z-Cy=`!RHFc#|1r}5%hRQ(Bm0Fzr7!q(Bm0Fk7oouo)PqTM$qFKL62ty{f>SfLXT$z zJ)RNtct+6U89|R{1U)`3=<$r8$1{Q+&j@-vBk1vrpvNv~$v1In! zxo;+q3m)s>`%J!{Ys=S!UZExQiY=j6Y6-n^OXyWvLa*8qdbO6&tG9$+qa}3zO}}1f z_THGgKep^P;r`jGCG^%Up=Vm+8-2adHglKezU?<++inqF|8*^)U*8gXyOz*zXbJts zme6l%3H|1l&~Ir8{nnPyZ)*v?eM{&aT0-yG5_+eW&^xz;-lZkLpU@Ke#Fo$}wS+#oCG;sRp-*iI zeOgQC(_2EH(GvR1me6Olgg(0^^f@h|&us~PUQ6inTS8yZ68gfH&=<9YzPKgyB`u*Z zZ3%r@OX$m6LSNAm`pTBjSG9z`x+V1DmeALV^mQ$vuWt!`LrdrzTSDK|68h$r z(6_XNzO^OvZ7rd1ZwY-zOXxdWLf_RA`tFv{_q2q*wkdgYeTtF(k(wI%dwEumL$3B5*3=>D6YooDZjx%*?w zZWHdGty)5F-4gn>EupuWyE*r5cl@^9B3yr6OX%0Pgx;l zT0+0ICG^``LT}#^dWV+KJGO-0sU`HzEunX53B7Ae=-paE@7@x6kCxDTwuIiRCG_4c zq4#MCy>CnC{aQlr-xB(Ome7k@LLb-?`k=qp-6 zU)d7+s+Q1Kw}f8Y68f5!(AT!azt-2agzMM0gubCA^o=c{Z)ypBb4%!3T0-C268g55 z(6_gQzN01doh_m7Y6*RJOXz!ALf_jG`o5OX54JNRrf22pS^1^j4hUz_Pqc*oWJ~Bz zwS@k3OX!zdLVu0(2uu-exfDxlP#g2 zva>bUPj{Pe{h5}~&$fhqt|j#IEumj%3H@SAO#hzhx90DuE;n~~zGlh%n#sQn^T7Pe ze4okp`$E4D%h!Znp(XT+EumLx3B7Vl=v7)mui6rNwU*GUw}f7!C3OG2@;Pa}W=-g| zT0*bg5_+AM(9YPCbG9DO+1eSK(BnBu63H`d3(64U^ygnnyF z=(n|m-o7RD4lSW~Yze(nOX!_jLhsTNde@fFyS0Sgy(RP>Eur^p3B6ZK=)GG)@6!@` z-TSA}L68ij>&=<6XzOW_qMJ=H(ZV7!!OXy2mLSNPr`tp|0SG0t_vL*CY zEupV&3B9-_^ffJ^uWbo^-Q3ODZ(o}4$;-Z2U*8qN_rIYf^o=c{Z)ypBb4%!3T0-C2 z68g55(6_gQzN01doh_m7Y6*RJOXz!ALf_jG`o5OX54JNRK0n`HhR@ai*b<+u|EVSP zKevSbmzL1)YYF}Sme3z)3H`yA&>w0E{o$6-A885w(U#C3YYF}Fme3Eignqas^dl{y zA8iT!SWD=~TS7n468g!O&`-65e!3;}GcBQ?Z3+EcOX%lYLch=w`o)&`O#QDd@#*@d zme8MQ3H`~I(4T4v{b@T}bNzC+3D-Z<68f_(p+DCW`tvQJzt9r;i!JeT{iT+0{mU(( zztR$)s{gGeK3RXYCG^)?LVvv_K2g7>za#O0iRI>X+sE|Fw}f7yCG?6dp;u}Ny>d(F zRa!!?+7f!Tme8xWgkGa1bpO3_cGhdwgkGy9^x7?<*J%m8ZcFI(T0*bi5_*G{&>Oad z-l!$?#x0>YX$ie)OX$s7LT}y@dW)9OTegJWswMQ+Eum*xV){AiWzW*?pPjp32)#>7 z=v`Yv@75A}_mZgkIDV`oNaZ2epJg zxFz%34K^g=)+q=AJG!}$d=GYwS+#pCG;^Zp^t3|eOycE<6Ax4eQ!(X`&vRj z*v^QUetpNXXJ?VEuo)l z3H^La=oeZ-zt|Gf&q*(PcK&An?ZoBh?#@@$UsyEz z*1zoYqR)Gd&x?9IBk1vrpq&p0{q~m7;~7DZX9PW-5%hRQ(Bm0Fk7oq^j(#3Ok7oou zo)PqTM$qG9fgT?V^!U7}$LB>oJ}>I=c~Otgi+X%s)Z_D_9-kNW_`Imc=S4j}FY57m zQIF4ydVF5gOS<0(Jw7k$@p(~?&x?9|UexdG$079iyr{?LMLj+*>hXC|kI##Gd|uS= z>gOT!_`Imc=S4j}FY0&qdW7~kFnKJ{hXC|kI##Gd|uS! z^P(P~7xnnOsK@6;Jw7k$@p(~uUYs#;|GaLyHhnBudPa7Ce5XGajBg7|dK)18+NX~T zufE5HMYF&A%RVmnv+(h8L62tyJ)RNtct+4~@5d$dct+6U89|R{1U;S+^msxO*2SVyEOalzxDmM?H1wnU)K`)^(~>dYYF{^ zme6l(3H_#)&~I)D{g#%{Z*2+vwwBP_w}jrICG?Ihp?7Kty>mciTS6bx68hMd(8sleKE5UN2`!;dYzci*OX!nZLZ8wS z`qY-tr?rGWy(RP+EuqhB34K;e=(AfwpVJch+?LSiwS+#uCG-U?p)YI+eNjv3i(5ip z(h~a8me7~Aguc8b^c5|kuWSi@RZHlrTS6~x34KjV=xbX-U)K`)`j*f)w1mE~CG<@# zp>J*peM?K|TU$ck))M;mme6;!gub&S^j$5X?`{cwPfO@~TSDL068gb*M#S{%;g|iI z9-l9Jr2B=?kG6zJcwuFAJCG_(xp_mUkLrRme5~siBHr%rhbo!<>qzU`_}R;p;u@Ly<$t~m0Cit+!A_~ zme8xVgkG&B^y)34*Jugdf77$`?7cB}e{9)p!u_*VOX#gzLcg{p^fq%h=f3TZ-?m$X z>#u7G{rZ;B+qHy#Lrds4wuF9DOXxSZgnmm)=(o0nep^fE?OQ_c&=Pvbme4!3gx;(0jFn-n%9AJ}sg5Z3(?!OX&SuLLbl)dQnU016x8L z)DrsOme7Z^gg&$-^kFTb4{r&5L`&!+TS6bz68h+t(8sidKDH(FaV?>bZwY-uOXw3@ zLZ8$U`s9|-r?iAVwI%dvEul|u34KOO=rdbFpVboj?3U2yw1hslCG>eMq0es#eL+j; z3tK{8)DrsQme7~9gub*T^kprfFK-EbMN8-_TS8ye68h?v(2H9_U(*u$+LrigeO*hq zetk>m8(Ko&*b@4tme4o1gubOE^sOzSZ)*vCdrRm$T0-C168f%|(08|lzNaPhy)B{d zYYF{eJ0oIxR-T@fU+V3E@b!nEXbJtvme8MS3H@n1TXX$#w+Yuj(-Qi#EulZx68iHk zp})`)`im{`?fSc;@16VUr)TBvhjr)g3y&@9b(_%Zw}jrHCG>_Zp*Ly?y>UzEO4 z+7f!Rme8BGgx;bh^p-85w`vK!bxY`(mY9B?d3ts}u(t!k+s8pIp$~2eeMn2_Lt8>0 z))M;gme5DEgg&w*^ieILk8TNlOiSovTS8CI&eOB=4$o|Abei-;g-;kw1j@NCG=x0p&xGv{X|RXCtE^4)e`#Yme9|%gnqUq^m8qtpKl5M zLQCiuTVnco=WosDogbLHJC`|2`+M5U_i{q7&=PvZme4D;gkHHN^eQc(S8WNsT1)8F zTSBkV61x9h`Mk4UvnKRfEuq(L3B68BXy@q3Sz3>0Y3&?M=wmWdOS<(nSLF_*Xq}{gzMYP-JJWj--vCyMY#UDme8+n z3B6rQ=r^>4eq&4MH?@R*b4%#Aw1j?ZOX#`uLX6C$xkq0ep!eNId0b6Y~6*An{tme3cpgubvP^hGV9 zFK!8aNlWNUTS8yf68iF%&{wpCzOp6sRV|^fZVA1(CG<5dp|5QTeO*iF>svzK&=UH_ zme4n~gub~Y^eru+Z*2*ETTAHMTSDK_68g@T(08?jzPly#JuRW{Z3%r}OXvsN84=Te zhh^Ebvwyebk?t2lKiU%dv6j$}w}gJ8CG?Xmp`Ws|HP=sfn{fS^me9|(gnq6i^z$vD zUuX&aVoOZ_9hPO!&fo06!?N<+-TBJ7Rr+^VR-b?Qq}N=~GYfk1e(uHp)a!fSobdJf zHGO@^Dw7wqGxN9q=dP_byxiB;jz6~^e{MZ_U9YAmj~V)$UU~ev|EI3|xgRjG+`Ml4 z{KWDtp;u@Ly<$t~m0Ci(|0egL_G_H{I`yhGq5WDXzh3RvI{EeLHEKfl-zy&@X!phB z{?KdHgkHNP^g1n}*KG;CUQ6ipTS9Np5_-dy&>OXc-nb?7CM}^iZ3(?uOX$s8LT}L$ zddrs3TeXDVx+U~XOH4oCzU;>c?w_5zUkJTROXyu&LhsfRdiR#ld$fe!vnBLiEur^r z3B6BC=zUv4@7EG~|CZ1Pw1i&N68gZF&LpU@Ke#Fo$}wS+#oCG;sRp-*iIeOgQC(_2EH(GvR1 zme6Olgg(0^^f@h|&us~PUQ6inTS8yZ68gfH&=<9YzPKgyB`u*ZZHfP-zN{r&zq}>% z6)mB!Yzci;OX#axLN9I!eN9W~YgFOH1fmTSDK~ z68iR*(08vvuP^&C!Z-VM-S3^d zJC`{-`}lkLUQXy0T0*bb5_+YU&?~ouUZo}Usx6^cYYDx2OXxLPLigV*XJ@@;P3W~+ zLa*HtdYzWg>$ZemuO;;QEulAP3B6%U=#5%JZ`=}kla|n%wuIiSCG_Skp|@xWy=6=2 zty)5F-4c4HC8qz*`Lbtc_s`DVFNEHuCG@T>p?7Nuy?aaOJz7HV*%ErMme6~*gx;qm z^u8^j_iG8ge@o~CT0$>s34LHo=!059AKVi9ke1MgwuC;cCG_Dfp^s<@ePm1Mqgp~A z-4gnkme9wxgg&k%^zkjBPiP5!VoT_gT0)=P68e;u(5JS9KCLD6=`EqpXbF90OX#y& zLZ96d`ka=~=eC4CuO;;PEuk-H34LKp=!;rHU)&PJ&oeOpWD z+gn24(GvR3me6;#guc5a^gS)1?`;WvUrXo*+Zhql&$lmomiGDSN4j4K{b)<*$67)^ z-V*wWme5bOgnr7-)?7c`ZNl|uT0%eD68gE8(9gGoexW7wi!CwzeEYIz>2LP)?aR&G zov&Fkzh?69oPVQhKEM5%K7T#AHhKT{QlDpc?Ew?M&*c00*mwDw&?~frUa=+gN-d#R zZVA0gOXyWwLa){mdi9pjYqW&!zgN!Zdd-^9Yqf-SHYfCWKG);tme9Mkgx0))M;gme5DEgg&w*^ieILk8TNlOiSovTS6b#68iX- z(9Y(Bb~Y#UNiCsIZV7!#OXyQuLZ8+W`t+92XS9SqvnBLdEuqhD34KmW=yO{_pVt!l z{Fcxcw1mE}CGO8CG;&Vp>J&oeOpWD+gn24(GvR3me6;#guc5a^gS)1 z?`;WvUrXo*+Zhql&uuSz9{0K6N4j4K{b)<*$67)^-V*wWme5bOgnp_e^wTY&pJ@sG zY)k0pT0%eH68eRf&@Z;c^mE(Gp2xr0&uu?2cX!U7o0prv?PJ~LyG`g7T0*bb5_+YU z&?~ouUZo}Usx6^cYYDx2OXxLPLigV*XJ@@;P3W~+La*HtdYzWg>$ZemuO;;QEulAP z3B6%U=#5%JZ`=}kla|n%wuIiSCG_Skp|@xWy=6=2ty)5F-4c4HC8nR-UiR$l{@JciTS6bx68hMd(8sleKE5UN2`!;d zYzci*OX!nZLZ8wS`qY-tr?rGWy(RP+EuqhB34K;e=(AfwpVJch+?LSiwS+#uCG-U? zp)YI+eNjv3i(5ip(h~a8miWKw%UZ(q%UeQU(GvQ~me5zVguc2Z^x~G#*R+Jbwk7m+ zEupV(3H^V!cc0x>-HUm})oau%NWd9!Zz_NUDm0jJvjwbU`+6j{97n!u4yK=q)Xwx3+}d))IPqOXwXfp?9`~-qjL%cRM9wyq7)m?Cjp}Yh4#Y zzupr1jh4`FwuFAGCG^`Zq2IC7HR*SIOnCmimeB9Fg#Mr<^oK2>KWYj6aZ8N%vS*&1 z|Ec$~|73DG$Fp^*nXXoc8A8w}4e$V%qa29-_CG?9epx2|b>j$NSE!CNGES*Td=_6J8(xH}{?W_0*p|@8wS+#tCG-g`p-*fH zeNs#4lUqWc(h~aAme8|WLeFjqeOgQC(_2E%X$d{ICG@=9eNRj1ds{-^*An{vme8wO zLa%NKy{0Af11+H+Yzh5POX!DNLO;?H2lavb(U$Q1V=bW{ZwdWGOXw$C;=q2rr&_}E zPfs3B|KCr(?>*o5&;7dH2lUTd->*#g=WS>Sy|E?qrk2o~TS9MX3B9!?^tP7J+gn2K zXbHWuCG@VA(7W3y5#xKWXP)D|SNXNB3!z_c3H?S(=r>zJzts}@?UvB*w1j@QCG>kO zq2F%_{Xt9U4_iWi)DrsRmKfiAJ@X#;_}=S%XKDXG{?WZg2z^XT=;2vf56{y2_SI zp@(N_Jv>Y68~bsDzNsbj%`KsCX$gI6OX%BLLf_sJ`i_>+ceaGSt0nZ^Eurse34L!% z==)ki-`^5?RZHmAEuq)6#KD~%hi7Te56{wic$U@=_wxunJWK20Sy~Uz(t3E7)=%{7 z5c3B9u=^sbiByW1%d<2~@1XXo+#BjY{rzv<5f;r%6l+Y+zU`~8xy`Qgvj0X<9T z16x8L)DrsOme7Z^gg&$-^kFTb4{r&5L`&!+TSE73?|o?9hx5Jnq4m)JT0&pg5_(BX=%p>8m$igm-V%B|JMViBhwE`;*M-nG zwS>O8CG;&Vp>J&oeOpWD+gn24(GvR3me6;#guc5a^gS)1?`;WvUrXrwTSBjD3B9@{ z^qQ8?Kc=&@f8YO?XXisZt8eT&5&rp`T0(Db3B9Ey^wyTp+gd_zZwbAlCG^gg(7Ref z?{23=jQ5?#`_7;AIv~cg^Z(h|dH>1f+V79QU*o;h?n&=?ZqJ`#-+P!Fyuajrwf9jS z(6fX-uqCv6)r212uhzr+)p~fpS`Y76>*4)sJ-lD7hxe;>-}c_G*2DYNdU(HD56|3s zc)wa7*RMzD<6A1Eun9134L2j=-XRD-_a8K&X&-3wS>OA zCG-8eP2uH`&&Y&=0hPey}AD?}PC}E#dixTS7n568h1W(2uo* ze!M006D^^iYzh5TOX%U5TR+qH6Z+Yf&}&;luWJdtz9saAme3npLT_pby}2dymX^?4 zTS9MZ3BA1~^p2L$J6l5UY6-o&of0v=PkQE=+xv=N>$(v7^_I|Yw1j@MCG=Y@q2F!^ z{Z32hcUwZg*An{ume3!xg#NH4^hYhBKW>TfebWDO@00dt>wvB?p$}{ceNap2gIhu$ z(h~a6me7Z_gg(3_^bswgk8BCuxAE*eeKjW6V@{6=*Jo}^=y@%nr&?k>JCA4Q6}=7! zuaApbLa%HIeQ`_ZOIku-+7kM*me7~CgubFB^p!24uWAW>bxUaHW$0LVs?jYtmoznDG3UEup_^3H^0T=xX(8sdz>^$Cgp1iVpjrDr?wSw^a`1O|1ztIx7*b;h4OX#I7p_jFUUfvRVJUfr~omWj>4%4rP)jcM> zKGw8^exN1vgDs&SY6<;tOXx>hLO0))M;gme5DEgg&w*bl>*gch-G4 z96!g<#R=DAc;8tM?>p*3j156{kecy`vqv$Gzao%Qnmb0_qQme3cqgkISa`r?+*m$Zbwv?cUqEuk-O z34KLN=qp=7U)2)&>Xy*gw1mF4CG>SIp|5WVeM3v=8(Tu()DrsUme9Afgub;U^ldGn zZ*K{GM@#5CTSDK}68i3z(D$^2zPBaxeJ!EyZwbAsCG_f+&}&*kKhP5T!IsbuwS<1S zCG;aLp&xAt{a8!r$6G=_(GvQ}me9kqvmTzE|J_kD&(5BIwx37nwJo97wS->Z5_&^R z=#4F*H?@S`+!A_AOX#gFp|`b!-rf>=M@#6PEunX{gx=jwi5TA}J@f25etnLaXJ_x5 z{*%JDzme60fg#M-_^tUacziWx{f8S+1JI|WD6sAA6&F(Sb zW8G;jp-*oKJ*Oq~+?LSuT0&2?#CUcd&(14)9S~k0&dr2g*%JEVme7~9gub*T^kprf zFK-EbMN8-_TS8ye68h?v(ATtt9?#C>+1dFSUJt{wvwo)UC-e__c0TH;@$CGQt}Edz z`qP%sKWhp7^On%hwS<1YCG-m|p)=(h~aBme8-Ygnqpx^cyXq-)sr} zR!iu&TSC8Mr)$#h_L%Vedo7{gZwdWDOXv?Lf_L8`rek% z_qBw+za{jlme8wPLa%9wWBa?~ftK+6gDs&SY6<;tOXx>hLOw1l495_(=s=&6<% z&(7o7c}1@S!t2Ainb0d+LXT(XU){5_|NMVUXXoSki{a^Bw}gNHXIes!XXnXz+V%Ll z>+zpk!ddc9TSEU?OXxpu3H_6n&_8Vn{j-+PKW_>BTubQZTSC9k68goK&@Z)wez_&| zD=nd4Z3+EaOX$~ILch@x`puTmZ?%MeyCw8HEur6S3H@G6==WPff6x;8!Y5_&v4kN2HdOZ5_&^R9N$@eV@r5`Q%mU0Eupuxgx=Z` zdRt5A?Jc2qw1nQ-5_(rl=-usi6`V7y`+PRs~!?Uv*3j156{kecy`vq zv$I~(zZapGwuD~R5_)+{=oKxYFKP+BvL*DzEuk-I34Lix=*wC{U)~b>ik8qsmrz-xB(Ume4o0gubaI^vx}yZ)pjAYfI?cT0-C668eso(08_k zzN;nl-7TT-X$gIAOX&MrLf_vKdR0s4)h(gdw1j@3CG>+Wp&x1q{cua@M_NKZ+7kM) zme7y4gnptW^ph>2pK6H{j+%LP_WbbdtcPc3?c7Z0buFQtn+fgQOlaq3LOVAT+PRs~ z&dr2&ZYH#IGohWE3B98w^v;&hyIMl;Zl^?y?}eUucJ_I9uXSAr{d!C2H(El!*%JD# zme6mvgnq|P*QDR=G2!|5T0+0y68eLd&>yyh{-`DN$1O2_9^uUU&g171?z`{oW5dzC zMhJaOOXy=;LLb)>`uLX6C$xkdOSPtd(Vcm<*Qv6LXT(X zefOQm&m)|C4&mQUULO0s*6ZV!eXgKCTL<(kp$}{ceNap2gIhu$(h~a6me7Z_gg(3_ z^bswgk8BCuxADI7^wpSLkMZoh?|(PN>*I{B4WZ9$34K;e=(Afw&uc99`8G^n!FsQUk|H$ zOn7~)X$kGz{Hu0$p8m~$|NCFx60XO1cHZ}XXJ^l!bX^Gj)0WUbYYF}Hme9|&gnqsy z^b0MaUu+5eQcLKUTSC9m5_&v4?|aY2_%%58+wX6BT@da||7}a`IXmw^$-V#mmOQEd zKl+!QtG!q1fSx7v@La8j=W0DXSL@-qS`W|FdU&qZ!*jJBo~!lnT&?@I_gt<0cBX$z z-~am_-4dQZrX}>TEuoKV34MG^=o4B(pV$)mq?XVpw}d{WCG@E+p=Y&(p4}4qw3g7P zw}c*^tM%NzpV0GKLQl2C_<4jg&(*GvGc}1Eun9134L2j=-XRD z-_a8K&X&-3wS>OACG-8eP2uH`&&Y(9YC^ey}CiLO<3L`tg>~Pqc)7vL*CWEuo)oiIb0-c~9E&&-VR$f7BBCLIurV5Eunwj68gE8(9gGoexW7wi!GsFY6<;vOXycxLciJ) zdOSPtdq2nLU5`(@E`(}jj zKZlQHM|4dHePm1MzK!>tr?1B3dN?q0ep!J-;RNf|k$= zTSA}H68hYh(2H6^FK!8aUQ6inTS8yZ68gfH&`VlEFKr3EtR?jFmeAwbdA#quYVvZJ zem$)2G2!*`Z|T0X-_JAsz6kwnOX#&Naq3ayedj;zIuOp5f7TNE&s#$Oq$Tutb{_9L z|09`w8uyGodeT34KXR=u2BdU)B=(@|Mt7 zw1mF0CG=G-p|5TUeN9W~YgFOH1fmTSDK~68iR* z(08*C&dr2g-4c3DOXvq$!v9|1gDv6thgw2E z+!Fecme7y3gnq0g^y4j|pJ)mFWJ~C$T0%eF68f2z(9gDnUfU9ST}$ZoEulBGgx=T^ zdQ(g2%`Ks~w1nQ;5_(%p=c9-uIr3@oRAWj}P=+`Cu`Wst9Z)yp>xh3?Lme5;ULT_scy}c#$ zj+W3nTSD(@2|b>j_uZ3rwtThgLg?}Ayzid0&)fU7>q6+yT0(zrr)$z*^qBDcmo1^c zY6<;yOXzP}LVw#5`n#4GKaX(Vd!hZ=KB8+v=p$P~_igWcp>-d=_P!Tdk7wuc?0kk7 zOL#q;*%JD!me6Ongr46LdO=I*g)O1aX$gI9OXx){p%=G=KCdP8`7NO@XbF8`OXwvn zp_jIVUe*$Nc}wW=>^$CgUNw0+OurtSn+dOvAM<@@|NcM3*?Dp`Cf8$qLHPB@v-9|Q zgl|q>4%27vw|Y!Ci@)6x`kj{0@3w?~uO;;REulYX3H@P9=#N@Lf7}w|edizMzVrT* z%enXOX?ePX=zjn3|IW9M69@Dxp$}{ceNap2gIhu$(h~a6me7Z_gg(3_^mwiw?@9me zU+ zef^3D`}u@l@1d5^54VJVq$TvDEukN43H^9W=qFl2KiLxcsg}?`#JPIiWN^~Q7c z`1Lv7`2O|a-0Dm{o~tJp#69CLdPwLmTS9-;68h_w(BHI#{5t8`drbJ)e_BiE(_2E%X$d{ICG@F zOH1fmTSDK~68iR*(08gnqFl^h+(FUv3HgN=xWhTSC9q68iO) z(9X@{*?IbMOrK>x?;+uA>)cG}FIz%?)e`#ameAj{g#NZA^mi@suXT1lWb%*Sd(T(8 zCw;{C-@m8(uDy7EYHv+`{J$Uc+@9Y(=sD9r`44yP|NZy<@!+|C?EbgU`#zu_OKA7L z2|c|3tq<<|3GMf{=l7@&tqFZtOX$N}Lc0!ou7y6bCUoESzOO|4txo@TzkmMdmhk*B zEuoKX34L5k=;K>LpU@Ke#Fo$}wS+#oCG;sRp-*iI?fkyy46lc0cs)GB>){z*56|#= zc!t-*GrXSapGW*dUnk_(`hLQX|Miy8ztIx7 z*b;h4OX#I7p_jFUUfvRVMN8<5T0*aE3GKX2=u28cU)mD-vX;=7w}if;CG?dop|5HQ zeRWIdYg$5I+YxSgkIYcdR=M@#6PEunX{gx=jwi5S1e-pn)n_;o@i z@5lbd +#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/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("= 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: + 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): + while True: + request = await self.ready.get() + 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) + + async def prepare(self, request, value): + try: + if request.canceled or self.closed: + return + loop = asyncio.get_running_loop() + 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): + while True: + value = await read_control(self.reader) + 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 == "submit": + 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_submit(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): + while True: + value = await read_frame(self.native_reader) + 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, + ) + ) + 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: + 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 + 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: + from tokenizers import Tokenizer + + 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") + tokenizer = Tokenizer.from_file(str(args.model / "tokenizer.json")) + 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" + ) + 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 + + @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): + text = self.preparer.prepare(value["text"], "en-US", value.get("apply_normalize", True)) + if not isinstance(text, str) or not text.strip(): + raise ValueError("text preparation returned empty text") + tokens = self.tokenizer.encode_batch([text], add_special_tokens=False)[0].ids + [131073] + 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("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("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 +747,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 +824,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..e9db3d1 100644 --- a/src/tts/magpietts/CMakeLists.txt +++ b/src/tts/magpietts/CMakeLists.txt @@ -13,10 +13,14 @@ add_library(nemo_speech_tts SHARED config.cpp ../c_api.cpp ../synthesizer.cpp + ../nemotron/native_runtime.cpp magpietts.cpp runtime.cpp ) +find_package(Threads REQUIRED) +target_link_libraries(nemo_speech_tts PRIVATE Threads::Threads) + # 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..f2ac495 --- /dev/null +++ b/src/tts/nemotron/native_runtime.cpp @@ -0,0 +1,544 @@ +// 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 "json.h" +#include "pcm_conversion.h" + +#ifndef _WIN32 +#include +#include +#include +#include +#include + +#include +#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(); +} + +} // namespace + +#ifndef _WIN32 + +struct NativeRuntime::Impl { + 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"); + 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); + } + 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(); + 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); + } + + 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) { + const std::string payload = "J" + value.dump(); + 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()) + 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); + } + + 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 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 (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 { + 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); + 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(); + pending.emplace(id, request); + } + Value submit = control("submit", id); + submit["text"] = prepared.metadata.original_text; + 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 { + 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; + } + + int fd = -1; + int timeout_seconds; + std::string name; + std::vector voices; + bool normalizes = false; + int width = 1; + std::string execution_mode; + 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/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..2abd29a 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,9 @@ 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) 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/nemotron_http_test.py b/tests/nemotron_http_test.py new file mode 100644 index 0000000..695d97a --- /dev/null +++ b/tests/nemotron_http_test.py @@ -0,0 +1,280 @@ +# 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 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): + with tempfile.TemporaryDirectory(prefix="nemo-native-http-") as temporary: + folder = Path(temporary) + worker = MockWorker(folder / "worker.sock", float_audio=float_audio) + 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", + ) + + 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) + arguments = parser.parse_args() + binary = arguments.binary.resolve(strict=True) + run(binary, runtime_seed=7) + run(binary, float_audio=True) + reject_runtime_options(binary) diff --git a/tests/nemotron_mock_worker.py b/tests/nemotron_mock_worker.py new file mode 100644 index 0000000..d3b64d0 --- /dev/null +++ b/tests/nemotron_mock_worker.py @@ -0,0 +1,172 @@ +# 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): + self.path = str(path) + self.float_audio = float_audio + 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, + ) + self.ready.set() + elif kind == "submit": + 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=3) + 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() From f3f3fb9458dac6fab18f7910ca9099d117e150b6 Mon Sep 17 00:00:00 2001 From: Aditya Chatterjee Date: Thu, 15 Oct 2026 10:18:36 +0000 Subject: [PATCH 2/3] Restore Nemotron B512 host-path performance --- THIRD_PARTY_NOTICES.md | 7 + docker/Dockerfile.nemotron | 20 +- docs/nemotron-integration-results.md | 269 ++++++ docs/nemotron-megakernel.md | 75 +- native/text_preparation/Cargo.lock | 771 ++++++++++++++++++ native/text_preparation/Cargo.toml | 37 + native/text_preparation/LICENSE.imported | 7 + native/text_preparation/NOTICE | 1 + .../crates/media-schema/Cargo.toml | 13 + .../crates/media-schema/src/lib.rs | 335 ++++++++ .../crates/processor-abi/Cargo.toml | 13 + .../crates/processor-abi/src/lib.rs | 678 +++++++++++++++ .../crates/processor-abi/src/raw.rs | 156 ++++ native/text_preparation/src/lib.rs | 464 +++++++++++ .../text_preparation/tests/qualification.py | 160 ++++ native/text_preparation/text_preparation.h | 18 + .../nemotron_cpp_bridge/native_preparation.py | 122 +++ python/nemotron_cpp_bridge/server.py | 209 ++++- scripts/nemotron_native/launch.py | 84 +- scripts/nemotron_native/tests/test_bridge.py | 261 +++++- .../tests/test_bridge_fairness.py | 180 ++++ scripts/nemotron_native/tests/test_launch.py | 241 ++++++ .../tests/test_native_preparation.py | 274 +++++++ server/http/http_server.cpp | 10 + src/tts/magpietts/CMakeLists.txt | 2 +- src/tts/nemotron/native_runtime.cpp | 124 ++- tests/nemotron_http_test.py | 78 +- tests/nemotron_mock_worker.py | 26 +- 28 files changed, 4579 insertions(+), 56 deletions(-) create mode 100644 native/text_preparation/Cargo.lock create mode 100644 native/text_preparation/Cargo.toml create mode 100644 native/text_preparation/LICENSE.imported create mode 100644 native/text_preparation/NOTICE create mode 100644 native/text_preparation/crates/media-schema/Cargo.toml create mode 100644 native/text_preparation/crates/media-schema/src/lib.rs create mode 100644 native/text_preparation/crates/processor-abi/Cargo.toml create mode 100644 native/text_preparation/crates/processor-abi/src/lib.rs create mode 100644 native/text_preparation/crates/processor-abi/src/raw.rs create mode 100644 native/text_preparation/src/lib.rs create mode 100644 native/text_preparation/tests/qualification.py create mode 100644 native/text_preparation/text_preparation.h create mode 100644 python/nemotron_cpp_bridge/native_preparation.py create mode 100644 scripts/nemotron_native/tests/test_bridge_fairness.py create mode 100644 scripts/nemotron_native/tests/test_launch.py create mode 100644 scripts/nemotron_native/tests/test_native_preparation.py diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 47cdf64..a950f5d 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -157,6 +157,13 @@ The legacy NVIDIA notice is preserved in [`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 diff --git a/docker/Dockerfile.nemotron b/docker/Dockerfile.nemotron index aa974d6..b375a09 100644 --- a/docker/Dockerfile.nemotron +++ b/docker/Dockerfile.nemotron @@ -2,6 +2,16 @@ # 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 @@ -52,6 +62,7 @@ RUN python3 /opt/nemotron-native/runtime_dependencies.py --remove-tensorrt --out 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 @@ -60,6 +71,8 @@ COPY LICENSE NOTICE THIRD_PARTY_NOTICES.md /opt/nemo-speech/share/licenses/nemo- 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 \ @@ -68,5 +81,10 @@ ENV PATH=/opt/nemo-speech/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/loc 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_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 index f0ae03e..d6a1034 100644 --- a/docs/nemotron-integration-results.md +++ b/docs/nemotron-integration-results.md @@ -225,3 +225,272 @@ 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`. diff --git a/docs/nemotron-megakernel.md b/docs/nemotron-megakernel.md index dc65e32..d8eb9f0 100644 --- a/docs/nemotron-megakernel.md +++ b/docs/nemotron-megakernel.md @@ -6,11 +6,14 @@ 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 bridge's tokenization and normalization. -The C++ frontend owns public HTTP handling and multiplexes requests over one -Unix socket. This integration is not a pure C++/CUDA rewrite of that controller. -There is no TensorRT-LLM execution dependency in the retained native path. +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 @@ -103,6 +106,32 @@ 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 @@ -123,6 +152,28 @@ 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. +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, @@ -185,13 +236,13 @@ 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. -To reproduce the optional float32-forwarding benchmark candidate, add -`-e NEMOTRON_BRIDGE_FORWARD_F32=1` before the image name in the `docker run` -command above. This moves float32-to-PCM16 conversion to the C++ frontend; -public HTTP audio remains PCM16. Leave the variable unset for the default -PCM16 bridge path, and record the selected value with benchmark receipts. -The [results report](nemotron-integration-results.md) distinguishes this -candidate from the default and records its qualification limits. +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: 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/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/server.py b/python/nemotron_cpp_bridge/server.py index 281a9b2..bdebbc4 100644 --- a/python/nemotron_cpp_bridge/server.py +++ b/python/nemotron_cpp_bridge/server.py @@ -8,6 +8,7 @@ import json import logging import os +import re import signal import time import uuid @@ -25,6 +26,33 @@ LOG = logging.getLogger(__name__) +def preparation_cpus(): + value = os.environ.get("NEMOTRON_BRIDGE_PREPARATION_CPUS") + if value is None: + return None + if not hasattr(os, "sched_getaffinity") or not hasattr(os, "sched_setaffinity"): + raise ValueError("preparation CPU affinity is unsupported on this platform") + available = os.sched_getaffinity(0) + selected = set() + for part in value.split(","): + match = re.fullmatch(r"([0-9]+)(?:-([0-9]+))?", part.strip()) + if match is None: + raise ValueError("invalid NEMOTRON_BRIDGE_PREPARATION_CPUS CPU list") + first = int(match[1]) + last = int(match[2]) if match[2] else first + if first > 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 @@ -40,6 +68,7 @@ class Request: 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) @@ -59,6 +88,7 @@ def __init__(self, bridge, reader, writer): self.output_bytes = 0 self.preparations = set() self.running_preparations = set() + self.dispatched_preparations = set() self.background = set() self.slabs = SlabReader() self.closed = False @@ -108,7 +138,11 @@ 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: + 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: @@ -148,8 +182,14 @@ def mark(self, request, event): 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 @@ -167,24 +207,51 @@ async def write_output(self): 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() - async with self.bridge.preparation_slots: - if request.canceled or self.closed: - return - request.preparation_started = True - self.running_preparations.add(request.preparation_task) + 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: - text, tokens = await loop.run_in_executor( - self.bridge.executor, self.bridge.prepare_text, value + result = await loop.run_in_executor( + self.bridge.executor, self.prepare_queued, request, value ) finally: - self.running_preparations.discard(request.preparation_task) - self.mark(request, "preparation_finished") + 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 @@ -210,13 +277,19 @@ async def prepare(self, request, value): 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 == "submit": + if kind in {"submit", "submit_prepared"}: if ident in self.requests: raise ValueError("request_id is already active") if ( @@ -229,7 +302,7 @@ async def read_client(self): self.requests[ident] = request self.native_requests[request.native_id] = request try: - self.bridge.validate_submit(value) + self.bridge.validate_input(value) except ValueError as error: await self.fail(request, "invalid_request", str(error), abort=False) continue @@ -245,8 +318,14 @@ async def read_client(self): 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") @@ -327,7 +406,11 @@ async def read_native(self): ) elif kind == "aborted": await self.fail( - request, "aborted", "native request aborted", abort=False, aborted=True + request, + "aborted", + "native request aborted", + abort=False, + aborted=True, ) else: raise ValueError("unsupported native output message") @@ -358,6 +441,11 @@ async def run(self): 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) @@ -377,7 +465,10 @@ async def run(self): finally: self.closed = True for task in (*tasks, *self.background): - if task not in self.running_preparations: + 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: @@ -393,6 +484,19 @@ async def run(self): 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) @@ -405,8 +509,6 @@ def __init__(self, args, *, tokenizer=None, preparer=None): ): raise ValueError("bridge requires the native streaming-v3 output contract") if tokenizer is None: - from tokenizers import Tokenizer - 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) @@ -419,7 +521,27 @@ def __init__(self, args, *, tokenizer=None, preparer=None): != tokenizer_seal["sha256"] ): raise ValueError("tokenizer does not match the selected sealed model") - tokenizer = Tokenizer.from_file(str(args.model / "tokenizer.json")) + 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 @@ -427,7 +549,10 @@ def __init__(self, args, *, tokenizer=None, preparer=None): self.tokenizer = tokenizer self.preparer = preparer self.executor = ThreadPoolExecutor( - max_workers=args.preparation_threads, thread_name_prefix="native-text" + 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 @@ -435,6 +560,23 @@ def __init__(self, args, *, tokenizer=None, preparer=None): 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 = { @@ -454,7 +596,13 @@ def validate_submit(value): 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", ""): + 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: @@ -463,10 +611,15 @@ def validate_submit(value): raise ValueError("apply_normalize must be boolean") def prepare_text(self, value): - text = self.preparer.prepare(value["text"], "en-US", value.get("apply_normalize", True)) + 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") - tokens = self.tokenizer.encode_batch([text], add_special_tokens=False)[0].ids + [131073] 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): @@ -518,7 +671,9 @@ def parse_args(argv=None): 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" + "--mode", + choices=("fixed_cohort", "fixed_request", "inflight"), + default="inflight", ) parser.add_argument("--width", type=int, required=True) parser.add_argument("--normalization-config", type=Path) @@ -560,7 +715,10 @@ async def serve(args): 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) + print( + json.dumps({"status": "BRIDGE_LISTENING", "socket": str(args.listen)}), + flush=True, + ) await stopping.wait() finally: if server is not None: @@ -573,7 +731,10 @@ async def serve(args): def main(argv=None): logging.basicConfig(level=logging.INFO) - print(json.dumps({"serialization": install_serialization(native_protocol)}), flush=True) + print( + json.dumps({"serialization": install_serialization(native_protocol)}), + flush=True, + ) asyncio.run(serve(parse_args(argv))) diff --git a/scripts/nemotron_native/launch.py b/scripts/nemotron_native/launch.py index 6a97afc..29da99a 100644 --- a/scripts/nemotron_native/launch.py +++ b/scripts/nemotron_native/launch.py @@ -5,6 +5,8 @@ import contextlib import json import os +import re +import shutil import signal import sys from pathlib import Path @@ -30,6 +32,8 @@ def parse_args(argv=None): parser.add_argument("--http-host", default="127.0.0.1") parser.add_argument("--http-threads", type=int) parser.add_argument("--frontend-binary", default="nemo-speech") + parser.add_argument("--native-cpus") + parser.add_argument("--frontend-cpus") parser.add_argument("--startup-timeout", type=float, default=900) parser.add_argument("--shutdown-timeout", type=float, default=30) args = parser.parse_args(argv) @@ -39,6 +43,10 @@ def parse_args(argv=None): parser.error("http-port must be in 1..65535") if args.http_threads is not None and args.http_threads <= 0: parser.error("http-threads must be positive") + try: + cpu_configuration(args) + except (OSError, ValueError) as error: + parser.error(str(error)) args.package = args.package.resolve() args.model = args.model.resolve() args.model_seal = (args.model_seal or args.model / "model-seal.json").resolve() @@ -50,6 +58,68 @@ def parse_args(argv=None): return args +def parse_cpus(value, available): + selected = set() + for part in value.split(","): + match = re.fullmatch(r"([0-9]+)(?:-([0-9]+))?", part.strip()) + if match is None: + raise ValueError("CPU sets must use comma-separated CPU IDs or inclusive ranges") + first = int(match[1]) + last = int(match[2]) if match[2] else first + if first > 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, "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}" @@ -146,7 +216,10 @@ async def observe(): async def run(args): + taskset, affinity = cpu_configuration(args) worker_command, bridge_command = commands(args) + worker_command = cpu_command(worker_command, "worker", taskset, affinity) + 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): @@ -178,7 +251,12 @@ async def run(args): processes = [] logs = [] tasks = [] - report = {"status": "STARTING", "width": args.width, "mode": args.mode} + report = { + "status": "STARTING", + "width": args.width, + "mode": args.mode, + "cpu_affinity": affinity, + } try: log = (args.receipts / "worker.log").open("wb") logs.append(log) @@ -192,6 +270,7 @@ async def run(args): await wait_ready( args.receipts / "worker/NATIVE_READY.json", processes, args.startup_timeout, stopping ) + record_cpu_affinity(affinity, "worker", worker) log = (args.receipts / "bridge.log").open("wb") logs.append(log) bridge = await asyncio.create_subprocess_exec( @@ -202,6 +281,7 @@ async def run(args): ) 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, @@ -218,6 +298,7 @@ async def run(args): ] 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( @@ -230,6 +311,7 @@ async def run(args): 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, diff --git a/scripts/nemotron_native/tests/test_bridge.py b/scripts/nemotron_native/tests/test_bridge.py index 3d29d34..630c181 100644 --- a/scripts/nemotron_native/tests/test_bridge.py +++ b/scripts/nemotron_native/tests/test_bridge.py @@ -11,6 +11,7 @@ import tempfile import threading import unittest +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from types import SimpleNamespace from unittest import mock @@ -25,11 +26,47 @@ pcm16_scalar, read_control, ) -from nemotron_cpp_bridge.server import Bridge, Request, Session, parse_args +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: @@ -165,6 +202,9 @@ def test_rejects_invalid_extent_type_identity_and_symlink(self): class BridgeTests(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): + environment = mock.patch.dict(os.environ, {"NEMOTRON_BRIDGE_FORWARD_F32": "0"}) + environment.start() + self.addCleanup(environment.stop) self.temp = tempfile.TemporaryDirectory() self.root = Path(self.temp.name) self.native_path = self.root / "worker.sock" @@ -189,6 +229,7 @@ async def asyncSetUp(self): ] ) self.bridge = Bridge(self.args, tokenizer=Tokenizer(), preparer=Preparer()) + self.bridge.frontend_preparation = getattr(self, "frontend_metadata", None) self.submitted = [] self.aborted = [] self.acknowledged = [] @@ -467,6 +508,16 @@ def blocked(value): self.assertFalse(self.bridge.preparation_slots.locked()) self.assertFalse(self.submitted) + async def test_prepared_submit_is_rejected_without_negotiation(self): + self.assertNotIn("text_preparation", self.ready) + self.writer.write( + control_frame("submit_prepared", request_id=1, text="hello", token_ids=[42, 131073]) + ) + event = await asyncio.wait_for(receive(self.reader), 3) + self.assertEqual((event["kind"], event["code"]), ("failed", "invalid_request")) + self.assertIn("negotiated mode", event["message"]) + self.assertFalse(self.submitted) + async def test_invalid_voice_fails_without_native_submission(self): self.writer.write(control_frame("submit", request_id=10, text="hello", voice="other")) event = await asyncio.wait_for(receive(self.reader), 3) @@ -501,6 +552,115 @@ async def test_tokenizer_identity_is_verified_before_loading(self): Bridge(self.args, preparer=Preparer()) self.assertEqual(factory.call_count, 1) + async def test_continuous_workers_advance_without_event_loop_permits(self): + self.bridge.continuous_preparation = True + self.bridge.executor.shutdown() + self.bridge.executor = ThreadPoolExecutor(max_workers=1) + first_started, release, second_started = ( + threading.Event(), + threading.Event(), + threading.Event(), + ) + original = self.bridge.prepare_text + + def work(value): + if value["text"] == "first": + first_started.set() + if not release.wait(5): + raise TimeoutError("first preparation was not released") + else: + second_started.set() + return original(value) + + with mock.patch.object(self.bridge, "prepare_text", side_effect=work): + try: + self.writer.write(control_frame("submit", request_id=1, text="first")) + await wait_until(first_started.is_set) + self.writer.write(control_frame("submit", request_id=2, text="second")) + await wait_until(lambda: len(self.bridge.session.dispatched_preparations) == 2) + release.set() + self.assertTrue(second_started.wait(1)) + events = [await asyncio.wait_for(receive(self.reader), 3) for _ in range(6)] + self.assertEqual(sum(event["kind"] == "complete" for event in events), 2) + finally: + release.set() + self.assertEqual({item["normalized_text"] for item in self.submitted}, {"first", "second"}) + + async def test_continuous_canceled_queue_is_bounded_and_skips_normalization(self): + self.bridge.continuous_preparation = True + self.bridge.executor.shutdown() + self.bridge.executor = ThreadPoolExecutor(max_workers=1) + self.args.max_pending_requests = 2 + session = self.bridge.session + started, release = threading.Event(), threading.Event() + original = self.bridge.prepare_text + + def work(value): + started.set() + if not release.wait(5): + raise TimeoutError("preparation was not released") + return original(value) + + with mock.patch.object(self.bridge, "prepare_text", side_effect=work) as prepare: + try: + self.writer.write(control_frame("submit", request_id=1, text="running")) + await wait_until(started.is_set) + self.writer.write(control_frame("submit", request_id=2, text="queued")) + await wait_until(lambda: len(session.dispatched_preparations) == 2) + self.writer.write(control_frame("abort", request_id=2)) + self.assertEqual( + (await asyncio.wait_for(receive(self.reader), 3))["kind"], "aborted" + ) + self.assertEqual(len(session.preparations), 2) + self.assertEqual(prepare.call_count, 1) + with self.assertLogs("nemotron_cpp_bridge.server", "ERROR") as errors: + self.writer.write(control_frame("submit", request_id=3, text="over limit")) + await wait_until(lambda: session.closed) + self.assertIn("client pending request limit exceeded", "\n".join(errors.output)) + self.assertEqual(len(session.dispatched_preparations), 2) + release.set() + self.assertEqual(await asyncio.wait_for(self.reader.read(), 3), b"") + self.assertFalse(session.dispatched_preparations) + self.assertFalse(session.preparations) + self.assertEqual(prepare.call_count, 1) + self.assertFalse(self.submitted) + finally: + release.set() + + async def test_continuous_shutdown_drains_active_preparation(self): + self.bridge.continuous_preparation = True + started, release = threading.Event(), threading.Event() + original = self.bridge.prepare_text + session = self.bridge.session + + def work(value): + started.set() + if not release.wait(5): + raise TimeoutError("preparation was not released") + return original(value) + + with mock.patch.object(self.bridge, "prepare_text", side_effect=work): + try: + self.writer.write(control_frame("submit", request_id=1, text="running")) + await wait_until(started.is_set) + self.writer.close() + await self.writer.wait_closed() + await wait_until(lambda: session.closed) + self.assertEqual(len(session.dispatched_preparations), 1) + self.assertFalse(self.worker_done.is_set()) + release.set() + await asyncio.wait_for(self.worker_done.wait(), 3) + self.assertFalse(session.dispatched_preparations) + self.assertFalse(session.preparations) + self.assertFalse(self.submitted) + finally: + release.set() + + async def test_continuous_configuration_rejects_ambiguous_value(self): + with mock.patch.dict(os.environ, {"NEMOTRON_BRIDGE_CONTINUOUS_PREPARATION": "true"}): + with self.assertRaisesRegex(ValueError, "must be 0 or 1"): + Bridge(self.args, tokenizer=Tokenizer(), preparer=Preparer()) + async def test_completion_accounting_failure_closes_connection(self): self.bad_completion = True self.writer.write(control_frame("submit", request_id=11, text="hello")) @@ -529,5 +689,104 @@ async def test_slow_request_queue_is_replaced_by_failure(self): self.assertEqual(session.output_bytes, len(request.output[0][0])) +class FrontendBridgeTests(unittest.IsolatedAsyncioTestCase): + frontend_metadata = { + "library": "/test/preparation.so", + "tokenizer": "/test/tokenizer.json", + "normalization_config": "null", + } + asyncSetUp = BridgeTests.asyncSetUp + asyncTearDown = BridgeTests.asyncTearDown + worker = BridgeTests.worker + + async def test_prepared_submission_preserves_tokens_and_skips_preparation(self): + self.assertEqual(self.ready["text_preparation"], self.frontend_metadata) + with mock.patch.object( + self.bridge, "prepare_text", side_effect=AssertionError("unexpected preparation") + ) as prepare: + self.writer.write( + control_frame( + "submit_prepared", + request_id=1, + text="Already normalized.", + token_ids=[42, 43, 131073], + seed=7, + ) + ) + events = [await asyncio.wait_for(receive(self.reader), 3) for _ in range(3)] + prepare.assert_not_called() + self.assertEqual([event["kind"] for event in events], ["accepted", "audio", "complete"]) + self.assertEqual(events[0]["processed_text"], "Already normalized.") + self.assertEqual(events[0]["token_count"], 3) + self.assertEqual(self.submitted[0]["normalized_text"], "Already normalized.") + self.assertEqual(self.submitted[0]["token_ids"], [42, 43, 131073]) + self.assertEqual(self.submitted[0]["seed"], 7) + + async def test_ordinary_submission_is_rejected_in_prepared_mode(self): + self.writer.write(control_frame("submit", request_id=1, text="hello")) + event = await asyncio.wait_for(receive(self.reader), 3) + self.assertEqual((event["kind"], event["code"]), ("failed", "invalid_request")) + self.assertIn("negotiated mode", event["message"]) + self.assertFalse(self.submitted) + self.writer.write( + control_frame("submit_prepared", request_id=2, text="recovery", token_ids=[42, 131073]) + ) + events = [await asyncio.wait_for(receive(self.reader), 3) for _ in range(3)] + self.assertEqual([event["kind"] for event in events], ["accepted", "audio", "complete"]) + + async def test_invalid_prepared_token_contract_is_rejected_before_dispatch(self): + for ident, tokens in enumerate( + ( + None, + [], + [1, 2], + [True, 131073], + [-1, 131073], + [131075, 131073], + [1.0, 131073], + ["1", 131073], + [131073] * 513, + {"id": 131073}, + ), + start=1, + ): + with self.subTest(tokens=str(tokens)[:60]): + self.writer.write( + control_frame( + "submit_prepared", request_id=ident, text="hello", token_ids=tokens + ) + ) + event = await asyncio.wait_for(receive(self.reader), 3) + self.assertEqual((event["kind"], event["code"]), ("failed", "invalid_request")) + self.assertIn("prepared token", event["message"]) + self.assertFalse(self.submitted) + self.bridge.validate_input( + {"kind": "submit_prepared", "text": "hello", "token_ids": [0] * 511 + [131073]} + ) + with self.assertRaisesRegex(ValueError, "prepared token"): + self.bridge.validate_input( + {"kind": "submit_prepared", "text": "hello", "token_ids": (42, 131073)} + ) + + async def test_prepared_request_cancellation_reaches_native_worker(self): + self.writer.write( + control_frame("submit_prepared", request_id=1, text="cancel", token_ids=[42, 131073]) + ) + self.assertEqual((await asyncio.wait_for(receive(self.reader), 3))["kind"], "accepted") + self.writer.write(control_frame("abort", request_id=1)) + self.assertEqual((await asyncio.wait_for(receive(self.reader), 3))["kind"], "aborted") + await wait_until(lambda: bool(self.aborted)) + self.assertEqual(self.aborted, [self.submitted[0]["request_id"]]) + + async def test_cancel_buffered_prepared_request_skips_native_submission(self): + self.writer.write( + control_frame("submit_prepared", request_id=1, text="hello", token_ids=[42, 131073]) + + control_frame("abort", request_id=1) + ) + self.assertEqual((await asyncio.wait_for(receive(self.reader), 3))["kind"], "aborted") + await wait_until(lambda: not self.bridge.session.preparations) + self.assertFalse(self.submitted) + + if __name__ == "__main__": unittest.main() diff --git a/scripts/nemotron_native/tests/test_bridge_fairness.py b/scripts/nemotron_native/tests/test_bridge_fairness.py new file mode 100644 index 0000000..79c63d4 --- /dev/null +++ b/scripts/nemotron_native/tests/test_bridge_fairness.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import asyncio +import json +import os +import sys +import unittest +from pathlib import Path +from types import MethodType, 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 +from nemotron_cpp_bridge.server import Bridge, Request, Session +from nemotron_native_backend.protocol import frame, message + + +class CountingReader(asyncio.StreamReader): + def __init__(self, payload): + super().__init__() + self.reads = 0 + self.feed_data(payload) + + async def readexactly(self, size): + value = await super().readexactly(size) + self.reads += 1 + return value + + @property + def frames(self): + return self.reads // 2 + + +class Writer: + def __init__(self, reader=None): + self.reader = reader + self.writes = [] + self.first_write = asyncio.Event() + self.input_frames_at_first_write = None + + def write(self, payload): + if not self.writes: + if self.reader is not None: + self.input_frames_at_first_write = self.reader.frames + self.first_write.set() + self.writes.append(payload) + + async def drain(self): + pass + + +def session(fair_batch, reader=None): + bridge = SimpleNamespace( + fair_batch=fair_batch, + diagnostics=None, + frontend_preparation=None, + validate_submit=Bridge.validate_submit, + args=SimpleNamespace( + max_pending_requests=1024, + request_queue_chunks=256, + max_output_bytes=1024 * 1024, + write_timeout=1, + ), + ) + bridge.validate_input = MethodType(Bridge.validate_input, bridge) + return Session(bridge, reader, Writer(reader)) + + +class FairnessConfigurationTests(unittest.IsolatedAsyncioTestCase): + async def test_default_boundaries_and_invalid_values(self): + plan = { + "physical_width": 1, + "host": {"batching_mode": "inflight"}, + "transport_compatibility": { + "output_contract": "pcm_f32_le_22050_mono_variable_streaming_v3" + }, + } + args = SimpleNamespace( + manifest=SimpleNamespace(read_bytes=lambda: json.dumps(plan).encode()), + width=1, + mode="inflight", + preparation_threads=1, + ) + preparer = SimpleNamespace(enabled=False, close=lambda: None) + for value in (None, "0", "64", "1024"): + environment = {} if value is None else {"NEMOTRON_BRIDGE_FAIR_BATCH": value} + with self.subTest(value=value), mock.patch.dict(os.environ, environment, clear=True): + bridge = Bridge(args, tokenizer=object(), preparer=preparer) + self.assertEqual(bridge.fair_batch, int(value or 0)) + await bridge.close() + for value in ("", "-1", "1025", "1.5", "true", "1_000", "999999999"): + with self.subTest(value=value), mock.patch.dict( + os.environ, {"NEMOTRON_BRIDGE_FAIR_BATCH": value}, clear=True + ), self.assertRaisesRegex(ValueError, "NEMOTRON_BRIDGE_FAIR_BATCH"): + Bridge(args, tokenizer=object(), preparer=preparer) + + +class FairnessTests(unittest.IsolatedAsyncioTestCase): + async def run_until_output(self, instance, consume): + tasks = [asyncio.create_task(consume()), asyncio.create_task(instance.write_output())] + try: + await asyncio.wait_for(instance.writer.first_write.wait(), 1) + for task in tasks: + if task.done(): + task.result() + return instance.writer.input_frames_at_first_write + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + instance.slabs.close() + + async def test_native_output_progresses_before_buffered_burst_finishes(self): + for fair_batch, expected in ((0, 192), (64, 64)): + with self.subTest(fair_batch=fair_batch): + reader = CountingReader( + b"".join(frame(message("accepted", request_id=str(i))) for i in range(192)) + ) + instance = session(fair_batch, reader) + instance.native_reader = reader + for i in range(192): + request = Request(i, str(i), "text", token_count=1) + instance.requests[i] = request + instance.native_requests[str(i)] = request + self.assertEqual( + await self.run_until_output(instance, instance.read_native), expected + ) + + async def test_invalid_client_frames_cannot_starve_ready_failures(self): + for fair_batch, expected in ((0, 192), (64, 64)): + with self.subTest(fair_batch=fair_batch): + reader = CountingReader( + b"".join(control_frame("submit", request_id=i, text="") for i in range(192)) + ) + instance = session(fair_batch, reader) + self.assertEqual( + await self.run_until_output(instance, instance.read_client), expected + ) + + async def test_retired_native_frames_still_yield_to_ready_output(self): + reader = CountingReader(frame(message("aborted", request_id="retired")) * 192) + instance = session(64, reader) + instance.native_reader = reader + instance.retired.append("retired") + instance.enqueue(Request(1, "ready", "text"), control_frame("accepted", request_id=1)) + self.assertEqual(await self.run_until_output(instance, instance.read_native), 64) + + async def test_output_yields_after_terminal_and_empty_queue_entries(self): + for empty in (False, True): + for fair_batch, expected in ((0, 192), (64, 64)): + with self.subTest(empty=empty, fair_batch=fair_batch): + instance = session(fair_batch) + payloads = [] + for i in range(192): + request = Request(i, str(i), "text") + if empty: + request.scheduled = True + instance.ready.put_nowait(request) + else: + payload = control_frame("complete", request_id=i) + payloads.append(payload) + instance.enqueue(request, payload, terminal=True) + task = asyncio.create_task(instance.write_output()) + try: + await asyncio.sleep(0) + if task.done(): + task.result() + processed = 192 - instance.ready.qsize() + self.assertEqual(processed, expected) + self.assertEqual(instance.writer.writes, payloads[:expected]) + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + instance.slabs.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/nemotron_native/tests/test_launch.py b/scripts/nemotron_native/tests/test_launch.py new file mode 100644 index 0000000..50e0b0a --- /dev/null +++ b/scripts/nemotron_native/tests/test_launch.py @@ -0,0 +1,241 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import asyncio +import importlib.util +import json +import os +import shutil +import signal +import socket +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +SPEC = importlib.util.spec_from_file_location( + "native_launch_affinity", Path(__file__).parents[1] / "launch.py" +) +launch = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(launch) + + +def options(native=None, frontend=None, port=None): + return SimpleNamespace(native_cpus=native, frontend_cpus=frontend, http_port=port) + + +class CpuConfigurationTests(unittest.TestCase): + def test_defaults_keep_commands_and_do_not_require_taskset_or_linux(self): + for platform in ("linux", "darwin"): + with self.subTest(platform=platform), mock.patch.object( + launch.sys, "platform", platform + ), mock.patch.object( + launch.os, "sched_getaffinity", return_value={12, 13}, create=True + ), mock.patch.object( + launch.shutil, "which" + ) as which: + taskset, configuration = launch.cpu_configuration(options(port=8000)) + self.assertIsNone(taskset) + which.assert_not_called() + for role in ("worker", "bridge", "frontend"): + command = ["program", "--option"] + self.assertIs( + launch.cpu_command(command, role, taskset, configuration), command + ) + self.assertIsNone(configuration["roles"][role]["requested_cpus"]) + + def test_cpu_ranges_are_validated_and_shared_sets_are_allowed(self): + with mock.patch.object(launch.sys, "platform", "linux"), mock.patch.object( + launch.os, "sched_getaffinity", return_value={12, 13, 36, 37}, create=True + ), mock.patch.object(launch.shutil, "which", return_value="/usr/bin/taskset"): + taskset, configuration = launch.cpu_configuration( + options("12-13, 36,12", "12-13,36-37", 8000) + ) + for role in ("worker", "bridge"): + self.assertEqual(configuration["roles"][role]["requested_cpus"], [12, 13, 36]) + self.assertEqual( + launch.cpu_command(["program"], role, taskset, configuration), + ["/usr/bin/taskset", "--cpu-list", "12,13,36", "program"], + ) + self.assertEqual(configuration["roles"]["frontend"]["requested_cpus"], [12, 13, 36, 37]) + for invalid in ( + "", + "-1", + "12,,13", + "13-12", + "12.0", + "12:13", + "14", + "12-36", + "999999999", + ): + with self.subTest(value=invalid), self.assertRaises(ValueError): + launch.cpu_configuration(options(native=invalid)) + + def test_each_control_only_prefixes_its_selected_roles(self): + with mock.patch.object(launch.sys, "platform", "linux"), mock.patch.object( + launch.os, "sched_getaffinity", return_value={12, 13}, create=True + ), mock.patch.object(launch.shutil, "which", return_value="/usr/bin/taskset"): + for args, selected in ( + (options(native="12", port=8000), {"worker", "bridge"}), + (options(frontend="13", port=8000), {"frontend"}), + ): + taskset, configuration = launch.cpu_configuration(args) + for role in ("worker", "bridge", "frontend"): + with self.subTest(role=role, selected=selected): + command = ["program"] + result = launch.cpu_command(command, role, taskset, configuration) + self.assertEqual(result != command, role in selected) + + def test_frontend_selection_requires_frontend_and_explicit_controls_require_tools(self): + with self.assertRaisesRegex(ValueError, "frontend-cpus requires http-port"): + launch.cpu_configuration(options(frontend="12")) + with mock.patch.object(launch.sys, "platform", "darwin"), self.assertRaisesRegex( + ValueError, "require Linux" + ): + launch.cpu_configuration(options(native="12")) + with mock.patch.object(launch.sys, "platform", "linux"), mock.patch.object( + launch.os, "sched_getaffinity", return_value={12}, create=True + ), mock.patch.object(launch.shutil, "which", return_value=None), self.assertRaisesRegex( + ValueError, "require taskset" + ): + launch.cpu_configuration(options(native="12")) + + +class CpuLaunchTests(unittest.IsolatedAsyncioTestCase): + async def test_unavailable_cpu_is_rejected_before_any_subprocess(self): + with mock.patch.object(launch.sys, "platform", "linux"), mock.patch.object( + launch.os, "sched_getaffinity", return_value={12}, create=True + ), mock.patch.object( + launch.shutil, "which", return_value="/usr/bin/taskset" + ), mock.patch.object( + launch.asyncio, "create_subprocess_exec" + ) as spawn: + with self.assertRaisesRegex(ValueError, "launcher CPU affinity"): + await launch.run(options(native="13")) + spawn.assert_not_called() + + @unittest.skipUnless( + sys.platform == "linux" and hasattr(os, "sched_getaffinity") and shutil.which("taskset"), + "Linux taskset required", + ) + async def test_role_affinity_precedes_child_threads_and_is_recorded_in_both_receipts(self): + available = sorted(os.sched_getaffinity(0)) + native_cpus = available[:2] + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + (root / "PACKAGE.json").write_text("{}") + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + port = listener.getsockname()[1] + args = launch.parse_args( + [ + "--package", + str(root), + "--model", + str(root), + "--width", + "1", + "--listen", + str(root / "bridge.sock"), + "--receipts", + str(root / "receipts"), + "--http-port", + str(port), + "--frontend-binary", + "/test/frontend", + "--native-cpus", + ",".join(map(str, native_cpus)), + "--frontend-cpus", + ",".join(map(str, available)), + "--startup-timeout", + "5", + "--shutdown-timeout", + "1", + ] + ) + + def child_code(path, listen=""): + return f"""import json,os,pathlib,socket,threading,time +observed={{'main':sorted(os.sched_getaffinity(0))}} +thread=threading.Thread(target=lambda: observed.update(thread=sorted(os.sched_getaffinity(0)))) +thread.start() +thread.join() +path=pathlib.Path({str(path)!r}) +path.parent.mkdir(parents=True,exist_ok=True) +path.write_text(json.dumps(observed)) +{listen} +time.sleep(60) +""" + + worker_path = args.receipts / "worker/NATIVE_READY.json" + bridge_path = root / "bridge-observed.json" + frontend_path = root / "frontend-observed.json" + worker_command = [sys.executable, "-c", child_code(worker_path)] + bridge_command = [ + sys.executable, + "-c", + child_code( + bridge_path, + f"s=socket.socket(socket.AF_UNIX); s.bind({str(args.listen)!r}); s.listen(1)", + ), + ] + frontend_code = child_code( + frontend_path, f"s=socket.socket(); s.bind(('127.0.0.1',{port})); s.listen(1)" + ) + original_spawn = asyncio.create_subprocess_exec + spawned = [] + + async def spawn(*command, **kwargs): + spawned.append(command) + if "/test/frontend" in command: + index = command.index("/test/frontend") + command = (*command[:index], sys.executable, "-c", frontend_code) + return await original_spawn(*command, **kwargs) + + callbacks = {} + loop = asyncio.get_running_loop() + with mock.patch.object( + launch, "commands", return_value=(worker_command, bridge_command) + ), mock.patch.object( + launch.asyncio, "create_subprocess_exec", side_effect=spawn + ), mock.patch.object( + loop, + "add_signal_handler", + side_effect=lambda number, callback: callbacks.setdefault(number, callback), + ): + task = asyncio.create_task(launch.run(args)) + try: + + async def ready(): + while not (args.receipts / "LAUNCH.json").exists(): + if task.done(): + task.result() + await asyncio.sleep(0.01) + + await asyncio.wait_for(ready(), 5) + finally: + if signal.SIGTERM in callbacks: + callbacks[signal.SIGTERM]() + await asyncio.wait_for(task, 5) + for role, path, expected in ( + ("worker", worker_path, native_cpus), + ("bridge", bridge_path, native_cpus), + ("frontend", frontend_path, available), + ): + self.assertEqual( + json.loads(path.read_text()), {"main": expected, "thread": expected} + ) + for name in ("LAUNCH.json", "EXIT.json"): + report = json.loads((args.receipts / name).read_text()) + record = report["cpu_affinity"]["roles"][role] + self.assertEqual(record["requested_cpus"], expected) + self.assertEqual(record["effective_cpus"], expected) + self.assertGreater(record["pid"], 0) + self.assertEqual(len(spawned), 3) + self.assertTrue(all(command[1] == "--cpu-list" for command in spawned)) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/nemotron_native/tests/test_native_preparation.py b/scripts/nemotron_native/tests/test_native_preparation.py new file mode 100644 index 0000000..ec61713 --- /dev/null +++ b/scripts/nemotron_native/tests/test_native_preparation.py @@ -0,0 +1,274 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import ctypes as C +import json +import os +import sys +import tempfile +import threading +import unittest +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from unittest import mock + +ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(ROOT / "python")) + +from nemotron_cpp_bridge.native_preparation import FrontendPreparation, NativePreparation + + +class Function: + def __init__(self, function): + self.function = function + + def __call__(self, *args): + return self.function(*args) + + +class Library: + def __init__(self): + self.buffers = {} + self.freed = [] + self.destroyed = [] + self.inputs = [] + self.result = {"text": "Hello", "tokens": [1, 2, 131073]} + self.error = None + self.status = 0 + self.started = None + self.release = None + self.nemotron_text_create = Function(self.create) + self.nemotron_text_prepare = Function(self.prepare) + self.nemotron_text_free = Function(self.free) + self.nemotron_text_destroy = Function(self.destroy) + + def assign(self, pointer, value): + C.cast(pointer, C.POINTER(C.c_void_p))[0] = C.c_void_p(value) + + def allocate(self, pointer, data): + buffer = C.create_string_buffer(data) + address = C.addressof(buffer) + self.buffers[address] = buffer + self.assign(pointer, address) + + def create(self, config, tokenizer, handle, error): + self.assign(handle, 123) + return 0 + + def prepare(self, handle, data, length, normalize, output, error): + self.inputs.append((C.string_at(data, length), bool(normalize))) + if self.started is not None: + self.started.set() + if self.release is not None and not self.release.wait(5): + raise TimeoutError("test did not release native preparation") + if self.result is not None: + data = ( + self.result if isinstance(self.result, bytes) else json.dumps(self.result).encode() + ) + self.allocate(output, data) + if self.error is not None: + self.allocate(error, self.error.encode()) + return self.status + + def free(self, pointer): + self.freed.append(pointer.value) + del self.buffers[pointer.value] + + def destroy(self, handle): + self.destroyed.append(handle.value) + + +class WrapperTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.library_path = Path(self.temp.name, "preparation.so") + self.library_path.touch() + self.library = Library() + self.patch = mock.patch( + "nemotron_cpp_bridge.native_preparation.C.CDLL", return_value=self.library + ) + self.patch.start() + self.addCleanup(self.patch.stop) + self.preparer = NativePreparation(self.library_path, "tokenizer.json") + self.addCleanup(self.preparer.close) + + def test_owned_output_is_freed_and_input_trimming_is_preserved(self): + self.assertEqual( + self.preparer.prepare_tokens(" \tHello\n", False), ("Hello", [1, 2, 131073]) + ) + self.assertEqual(self.library.inputs, [(b"Hello", False)]) + self.assertFalse(self.library.buffers) + self.assertEqual(len(self.library.freed), 1) + + def test_bad_request_types_do_not_call_native_code(self): + for text, normalize in ((None, True), (123, True), ("hello", "yes"), ("hello", 1)): + with self.subTest(text=text, normalize=normalize), self.assertRaises(ValueError): + self.preparer.prepare_tokens(text, normalize) + self.assertFalse(self.library.inputs) + + def test_native_errors_and_malformed_results_release_every_allocation(self): + for status, exception in ((1, ValueError), (2, RuntimeError)): + self.library.status, self.library.error = status, "rejected" + with self.subTest(status=status), self.assertRaisesRegex(exception, "rejected"): + self.preparer.prepare_tokens("hello") + self.assertFalse(self.library.buffers) + self.assertEqual(self.preparer._active, 0) + self.library.status, self.library.error, self.library.result = 0, None, b"not json" + with self.assertRaises(ValueError): + self.preparer.prepare_tokens("hello") + self.assertFalse(self.library.buffers) + self.assertEqual(self.preparer._active, 0) + self.library.result = None + with self.assertRaisesRegex(RuntimeError, "empty result"): + self.preparer.prepare_tokens("hello") + + def test_cancel_and_close_preserve_running_call_ownership(self): + self.library.started, self.library.release = threading.Event(), threading.Event() + with ThreadPoolExecutor(max_workers=1) as workers, ThreadPoolExecutor( + max_workers=1 + ) as closer: + running = workers.submit(self.preparer.prepare_tokens, "running") + self.assertTrue(self.library.started.wait(3)) + queued = workers.submit(self.preparer.prepare_tokens, "queued") + self.assertTrue(queued.cancel()) + self.assertFalse(running.cancel()) + closed = closer.submit(self.preparer.close) + try: + with self.preparer._condition: + self.assertTrue( + self.preparer._condition.wait_for(lambda: self.preparer._closed, 3) + ) + self.assertFalse(closed.done()) + self.assertFalse(self.library.destroyed) + with self.assertRaisesRegex(RuntimeError, "closed"): + self.preparer.prepare_tokens("late") + finally: + self.library.release.set() + self.assertEqual(running.result(3), ("Hello", [1, 2, 131073])) + closed.result(3) + self.preparer.close() + self.assertEqual(self.library.destroyed, [123]) + self.assertEqual(self.library.inputs, [(b"running", True)]) + self.assertFalse(self.library.buffers) + + def test_invalid_configuration_is_rejected(self): + config = Path(self.temp.name, "config.json") + for wrapper in (NativePreparation, FrontendPreparation): + for value in (None, False, [], "wrong", 1): + config.write_text(json.dumps(value)) + with self.subTest(wrapper=wrapper.__name__, value=value), self.assertRaisesRegex( + ValueError, "normalization configuration must be an object" + ): + wrapper(self.library_path, "tokenizer.json", config) + + def test_absent_configuration_and_explicit_disabled_configuration_remain_supported(self): + config = Path(self.temp.name, "config.json") + config.write_text(json.dumps({"enabled": False})) + for wrapper in (NativePreparation, FrontendPreparation): + for path in (None, config): + with self.subTest(wrapper=wrapper.__name__, config=path): + preparer = wrapper(self.library_path, "tokenizer.json", path) + try: + self.assertIs(preparer.enabled, False) + if isinstance(preparer, FrontendPreparation): + self.assertEqual( + json.loads(preparer.metadata["normalization_config"]), + None if path is None else {"enabled": False}, + ) + finally: + preparer.close() + + +NATIVE_LIBRARY = os.environ.get("NEMOTRON_NATIVE_PREPARATION_TEST_LIBRARY") +TOKENIZER = os.environ.get("NEMOTRON_NATIVE_PREPARATION_TEST_TOKENIZER") +NORMALIZATION = os.environ.get("NEMOTRON_NORMALIZATION_CONFIG") +CORPUS = os.environ.get("NEMOTRON_NATIVE_PREPARATION_TEST_CORPUS") + + +@unittest.skipUnless(NATIVE_LIBRARY and TOKENIZER and NORMALIZATION, "native fixtures not provided") +class NativeParityTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + from nemotron_cpp_bridge.normalization import TextPreparer + from tokenizers import Tokenizer + + cls.reference = TextPreparer(NORMALIZATION) + cls.tokenizer = Tokenizer.from_file(TOKENIZER) + cls.native = NativePreparation(NATIVE_LIBRARY, TOKENIZER, NORMALIZATION) + cls.texts = [ + " Hello world ", + "I paid $12.50 for 3 books.", + "Call me at 5 p.m.", + "Dr. Smith can't arrive until 3:30 p.m.", + "The year is 2026.", + "Call 1-800-555-0100.", + "Cafe\u0301 café naïve façade", + "中文 日本語 한국어", + "😀🧑🏽‍🚀", + "line one\nline two\tthree\r\nfour", + "\x1c\x1d\x1e\x1f\u0085\u00a0Hello!\u00a0\x1f", + "Keep spaces , and punctuation !", + "", + " \t\n", + "word " * 513, + ] + if CORPUS: + cls.texts.extend( + line.split("\t", 1)[1] for line in Path(CORPUS).read_text().splitlines() + ) + + @classmethod + def tearDownClass(cls): + cls.native.close() + cls.reference.close() + + @classmethod + def expected(cls, text, normalize): + normalized = cls.reference.prepare(text, "en-US", normalize) + tokens = cls.tokenizer.encode_batch([normalized], add_special_tokens=False)[0].ids + [ + 131073 + ] + if not 1 <= len(tokens) <= 512 or any(not 0 <= token < 131075 for token in tokens): + raise ValueError("token budget or vocabulary") + return normalized, tokens + + def test_exact_corpus_and_edge_cases_with_and_without_normalization(self): + for text in self.texts: + for normalize in (False, True): + with self.subTest(text=text[:60], normalize=normalize): + try: + expected = self.expected(text, normalize) + except (ValueError, RuntimeError) as error: + with self.assertRaises(type(error)): + self.native.prepare_tokens(text, normalize) + else: + self.assertEqual(self.native.prepare_tokens(text, normalize), expected) + + def test_parallel_preparation_matches_reference(self): + texts = self.texts[:5] * 26 + expected = [self.expected(text, True) for text in texts] + with ThreadPoolExecutor(max_workers=32) as executor: + actual = list(executor.map(self.native.prepare_tokens, texts)) + self.assertEqual(actual, expected) + + def test_invalid_native_configuration_is_rejected(self): + base = json.loads(Path(NORMALIZATION).read_text()) + with tempfile.TemporaryDirectory() as temporary: + path = Path(temporary, "config.json") + for change in ( + {"unknown": True}, + {"enabled": "yes"}, + {"pool_size": 0}, + {"pool_size": -1}, + {"language_code": "de-DE"}, + {"max_input_bytes": 0}, + {"normalizer_config": "../outside"}, + {"normalizer_root": base["native_library_root"]}, + ): + path.write_text(json.dumps({**base, **change})) + with self.subTest(change=change), self.assertRaises(ValueError): + NativePreparation(NATIVE_LIBRARY, TOKENIZER, path) + + +if __name__ == "__main__": + unittest.main() diff --git a/server/http/http_server.cpp b/server/http/http_server.cpp index 8c562dc..425865d 100644 --- a/server/http/http_server.cpp +++ b/server/http/http_server.cpp @@ -13,6 +13,7 @@ #include #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); diff --git a/src/tts/magpietts/CMakeLists.txt b/src/tts/magpietts/CMakeLists.txt index e9db3d1..21b6329 100644 --- a/src/tts/magpietts/CMakeLists.txt +++ b/src/tts/magpietts/CMakeLists.txt @@ -19,7 +19,7 @@ add_library(nemo_speech_tts SHARED ) find_package(Threads REQUIRED) -target_link_libraries(nemo_speech_tts PRIVATE Threads::Threads) +target_link_libraries(nemo_speech_tts PRIVATE Threads::Threads ${CMAKE_DL_LIBS}) # 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++ diff --git a/src/tts/nemotron/native_runtime.cpp b/src/tts/nemotron/native_runtime.cpp index f2ac495..1fc2526 100644 --- a/src/tts/nemotron/native_runtime.cpp +++ b/src/tts/nemotron/native_runtime.cpp @@ -19,6 +19,7 @@ #include "pcm_conversion.h" #ifndef _WIN32 +#include #include #include #include @@ -68,6 +69,112 @@ 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 @@ -146,6 +253,8 @@ struct NativeRuntime::Impl { 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 (...) { @@ -381,6 +490,13 @@ struct NativeRuntime::Impl { 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; { @@ -394,8 +510,11 @@ struct NativeRuntime::Impl { request = std::make_shared(); pending.emplace(id, request); } - Value submit = control("submit", id); - submit["text"] = prepared.metadata.original_text; + 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; @@ -479,6 +598,7 @@ struct NativeRuntime::Impl { 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; diff --git a/tests/nemotron_http_test.py b/tests/nemotron_http_test.py index 695d97a..b105ed1 100644 --- a/tests/nemotron_http_test.py +++ b/tests/nemotron_http_test.py @@ -10,6 +10,7 @@ import socket import struct import subprocess +import sys import tempfile import time import wave @@ -77,10 +78,12 @@ def failed_stream(port, text, require_no_audio=False): connection.close() -def run(binary, float_audio=False, runtime_seed=None): +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) + 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( @@ -137,13 +140,22 @@ def run(binary, float_audio=False, runtime_seed=None): ) first = response.read(512) require(len(first) == 512, "missing first PCM chunk") - require(not worker.completed, "worker completed before first audio arrived") + 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") + 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") + require( + first + rest == expected, + "F32 conversion changed PCM16 bytes", + ) finally: worker.release_stream.set() response.close() @@ -155,7 +167,8 @@ def run(binary, float_audio=False, runtime_seed=None): ] results = [future.result(timeout=10) for future in futures] require( - worker.parallel_ready.is_set(), "requests did not overlap inside the worker" + 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), @@ -178,10 +191,14 @@ def run(binary, float_audio=False, runtime_seed=None): 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") + 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" + speech(port, "recovery")[0] == 200, + "request failure poisoned other requests", ) connection, response = open_speech(port, "cancel") @@ -191,7 +208,10 @@ def run(binary, float_audio=False, runtime_seed=None): ) response.close() connection.close() - require(worker.aborted.wait(5), "HTTP disconnect did not abort backend request") + require( + worker.aborted.wait(5), + "HTTP disconnect did not abort backend request", + ) require( speech(port, "after-cancel")[0] == 200, "cancellation poisoned subsequent requests", @@ -201,6 +221,35 @@ def run(binary, float_audio=False, runtime_seed=None): "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( @@ -273,8 +322,19 @@ def reject_runtime_options(binary): 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 index d3b64d0..4e12ba8 100644 --- a/tests/nemotron_mock_worker.py +++ b/tests/nemotron_mock_worker.py @@ -8,9 +8,10 @@ class MockWorker: - def __init__(self, path, float_audio=False): + 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() @@ -78,9 +79,12 @@ def _run(self): batch_size=32, mode="inflight", text_normalization_enabled=False, + **({"text_preparation": self.preparation} if self.preparation else {}), ) self.ready.set() - elif kind == "submit": + 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: @@ -113,10 +117,17 @@ def _synthesize(self, message, abort): try: if text == "failure": self._json( - kind="failed", request_id=request_id, message="injected preparation failure" + kind="failed", + request_id=request_id, + message="injected preparation failure", ) return - self._json(kind="accepted", request_id=request_id, processed_text=text, token_count=3) + 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 = ( @@ -142,7 +153,12 @@ def _synthesize(self, message, abort): 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) + 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): From 721c67db94b42658b714faa739c24f72c3ecf2da Mon Sep 17 00:00:00 2001 From: Aditya Chatterjee Date: Thu, 15 Oct 2026 11:22:59 +0000 Subject: [PATCH 3/3] --- .clang-format | 1 + docs/nemotron-integration-results.md | 103 +++++- docs/nemotron-megakernel.md | 65 ++++ scripts/nemotron_native/launch.py | 112 +++++- .../tests/test_launch_direct.py | 336 ++++++++++++++++++ src/tts/magpietts/CMakeLists.txt | 2 + src/tts/nemotron/native_runtime.cpp | 333 +++++++++++++++-- src/tts/nemotron/native_worker_protocol.cpp | 260 ++++++++++++++ src/tts/nemotron/native_worker_protocol.h | 45 +++ tests/cpp/CMakeLists.txt | 10 + .../common/test_nemotron_worker_protocol.cpp | 300 ++++++++++++++++ 11 files changed, 1513 insertions(+), 54 deletions(-) create mode 100644 scripts/nemotron_native/tests/test_launch_direct.py create mode 100644 src/tts/nemotron/native_worker_protocol.cpp create mode 100644 src/tts/nemotron/native_worker_protocol.h create mode 100644 tests/cpp/common/test_nemotron_worker_protocol.cpp 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/docs/nemotron-integration-results.md b/docs/nemotron-integration-results.md index d6a1034..212212f 100644 --- a/docs/nemotron-integration-results.md +++ b/docs/nemotron-integration-results.md @@ -2,10 +2,10 @@ 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. B512 -startup latency remains higher than the legacy frontend in the same-host -control. These samples do not establish production readiness, acoustic -quality, or a performance improvement over the legacy deployment. +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 @@ -494,3 +494,98 @@ 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 index d8eb9f0..13931d3 100644 --- a/docs/nemotron-megakernel.md +++ b/docs/nemotron-megakernel.md @@ -152,6 +152,16 @@ 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 @@ -358,6 +368,61 @@ four refill cohorts after two warmup rounds: --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 diff --git a/scripts/nemotron_native/launch.py b/scripts/nemotron_native/launch.py index 29da99a..e81347a 100644 --- a/scripts/nemotron_native/launch.py +++ b/scripts/nemotron_native/launch.py @@ -3,6 +3,7 @@ import argparse import asyncio import contextlib +import hashlib import json import os import re @@ -32,6 +33,7 @@ def parse_args(argv=None): parser.add_argument("--http-host", default="127.0.0.1") parser.add_argument("--http-threads", type=int) parser.add_argument("--frontend-binary", default="nemo-speech") + parser.add_argument("--direct-worker", action="store_true") parser.add_argument("--native-cpus") parser.add_argument("--frontend-cpus") parser.add_argument("--startup-timeout", type=float, default=900) @@ -41,6 +43,8 @@ def parse_args(argv=None): parser.error("timeouts must be positive") if args.http_port is not None and not 1 <= args.http_port <= 65535: parser.error("http-port must be in 1..65535") + if args.direct_worker and args.http_port is None: + parser.error("direct-worker requires http-port") if args.http_threads is not None and args.http_threads <= 0: parser.error("http-threads must be positive") try: @@ -87,7 +91,9 @@ def cpu_configuration(args): 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, "bridge": native} + roles = {"worker": native} + if not getattr(args, "direct_worker", False): + roles["bridge"] = native if args.http_port is not None: roles["frontend"] = frontend configuration = { @@ -183,6 +189,66 @@ def commands(args): 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(): @@ -219,7 +285,10 @@ async def run(args): taskset, affinity = cpu_configuration(args) worker_command, bridge_command = commands(args) worker_command = cpu_command(worker_command, "worker", taskset, affinity) - bridge_command = cpu_command(bridge_command, "bridge", 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): @@ -234,6 +303,7 @@ async def run(args): 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 ""), @@ -248,6 +318,12 @@ async def run(args): 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 = [] @@ -256,6 +332,7 @@ async def run(args): "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") @@ -271,17 +348,18 @@ async def run(args): args.receipts / "worker/NATIVE_READY.json", processes, args.startup_timeout, stopping ) record_cpu_affinity(affinity, "worker", worker) - 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 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, @@ -289,7 +367,7 @@ async def run(args): "--tts-backend", "nemotron", "--native-socket", - str(args.listen), + str(args.native_worker_socket if direct else args.listen), "--port", str(args.http_port), "--host", @@ -303,7 +381,7 @@ async def run(args): logs.append(log) frontend = await asyncio.create_subprocess_exec( *frontend_command, - env=environment, + env=frontend_environment, stdout=log, stderr=asyncio.subprocess.STDOUT, ) @@ -318,7 +396,9 @@ async def run(args): "host": args.http_host, "port": args.http_port, } - report.update(status="READY", socket=str(args.listen)) + 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()) diff --git a/scripts/nemotron_native/tests/test_launch_direct.py b/scripts/nemotron_native/tests/test_launch_direct.py new file mode 100644 index 0000000..fd27767 --- /dev/null +++ b/scripts/nemotron_native/tests/test_launch_direct.py @@ -0,0 +1,336 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import asyncio +import hashlib +import importlib.util +import json +import os +import signal +import socket +import stat +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(ROOT / "python")) +SPEC = importlib.util.spec_from_file_location( + "native_launch_direct", Path(__file__).parents[1] / "launch.py" +) +launch = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(launch) + + +class Fixture: + def __init__(self, root, direct=True, width=512): + self.root = root + self.package = root / 'package' + self.model = root / 'model' + self.model.mkdir() + self.tokenizer = self.model / 'tokenizer.json' + self.tokenizer.write_text('{"fixture":true}') + self.library = root / 'libpreparation.so' + self.library.touch() + self.serving = self.package / 'programs' / ('b' + str(width)) / 'serving.json' + self.serving.parent.mkdir(parents=True) + self.plan = { + 'physical_width': width, + 'host': {'batching_mode': 'inflight'}, + 'model': {'checkpoint_sha256': 'a' * 64}, + 'transport_compatibility': { + 'output_contract': 'pcm_f32_le_22050_mono_variable_streaming_v3', + 'installed_transport_versions': { + 'trt_version': 'not-installed', + 'trtllm_version': 'not-installed', + 'trtllm_commit': 'not-used', + }, + }, + } + self.save_plan() + self.seal = { + 'model_identity': 'a' * 64, + 'files': { + 'tokenizer.json': { + 'bytes': self.tokenizer.stat().st_size, + 'sha256': hashlib.sha256(self.tokenizer.read_bytes()).hexdigest(), + } + }, + } + self.seal_path = self.model / 'model-seal.json' + self.save_seal() + (self.package / 'PACKAGE.json').write_text('{}') + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + self.arguments = [ + '--package', + str(self.package), + '--model', + str(self.model), + '--width', + str(width), + '--listen', + str(root / 'bridge.sock'), + '--receipts', + str(root / 'receipts'), + '--http-port', + str(port), + '--frontend-binary', + '/test/frontend', + '--startup-timeout', + '3', + '--shutdown-timeout', + '0.05', + ] + if direct: + self.arguments.append('--direct-worker') + self.args = launch.parse_args(self.arguments) + + def save_plan(self): + self.serving.write_text(json.dumps(self.plan)) + + def save_seal(self): + self.seal_path.write_text(json.dumps(self.seal)) + + def configuration(self): + with mock.patch.dict(os.environ, NEMOTRON_NATIVE_PREPARATION_LIBRARY=str(self.library)): + return launch.direct_worker_configuration(self.args) + + +class MetadataTests(unittest.TestCase): + def test_default_path_is_unchanged_and_direct_requires_http(self): + with tempfile.TemporaryDirectory() as temporary: + fixture = Fixture(Path(temporary), direct=False) + self.assertFalse(fixture.args.direct_worker) + args = fixture.arguments[:] + port = args.index('--http-port') + del args[port : port + 2] + with self.assertRaises(SystemExit): + launch.parse_args(args + ['--direct-worker']) + worker, bridge = launch.commands(fixture.args) + self.assertIn('--qualification-only', worker) + self.assertIn('nemotron_cpp_bridge', bridge) + self.assertNotIn('DIRECT_WORKER', str(worker)) + + def test_metadata_binds_exact_manifest_model_tokenizer_width_versions_and_normalization(self): + for width in (1, 32, 512, 4096): + with self.subTest(width=width), tempfile.TemporaryDirectory() as temporary: + fixture = Fixture(Path(temporary), width=width) + config = fixture.configuration() + self.assertEqual( + config['manifest_sha256'], + hashlib.sha256(fixture.serving.read_bytes()).hexdigest(), + ) + self.assertEqual(config['model_sha256'], 'a' * 64) + self.assertEqual(config['enabled_buckets'], sorted({1, width})) + self.assertEqual( + config['installed_transport_versions'], + fixture.plan['transport_compatibility']['installed_transport_versions'], + ) + self.assertEqual(config['mode'], 'inflight') + self.assertEqual(config['text_preparation']['tokenizer'], str(fixture.tokenizer)) + self.assertEqual(config['text_preparation']['library'], str(fixture.library)) + self.assertFalse(config['text_normalization_enabled']) + normalization = fixture.root / 'normalization.json' + normalization.write_text('{"enabled":true,"workers":16}') + fixture.args.normalization_config = normalization + config = fixture.configuration() + self.assertTrue(config['text_normalization_enabled']) + self.assertEqual( + json.loads(config['text_preparation']['normalization_config']), + {'enabled': True, 'workers': 16}, + ) + + def test_width_mode_output_and_model_identity_rejected(self): + changes = [ + ('physical_width', 32), + ('physical_width', True), + ('host', {'batching_mode': 'fixed_cohort'}), + ('model', {'checkpoint_sha256': 'invalid'}), + ('model', {'checkpoint_sha256': 123}), + ] + for key, value in changes: + with self.subTest(key=key, value=value), tempfile.TemporaryDirectory() as temporary: + fixture = Fixture(Path(temporary)) + fixture.plan[key] = value + fixture.save_plan() + with self.assertRaises(ValueError): + fixture.configuration() + with tempfile.TemporaryDirectory() as temporary: + fixture = Fixture(Path(temporary)) + fixture.plan['transport_compatibility']['output_contract'] = 'pcm16' + fixture.save_plan() + with self.assertRaisesRegex(ValueError, 'streaming-v3'): + fixture.configuration() + + def test_malformed_transport_versions_rejected_before_launch(self): + bad_versions = [ + None, + [], + {}, + {'trt_version': 'x'}, + {'trt_version': True, 'trtllm_version': 'x', 'trtllm_commit': 'x'}, + {'trt_version': '', 'trtllm_version': 'x', 'trtllm_commit': 'x'}, + {'trt_version': 'x' * 1025, 'trtllm_version': 'x', 'trtllm_commit': 'x'}, + ] + for value in bad_versions: + with self.subTest(value=value), tempfile.TemporaryDirectory() as temporary: + fixture = Fixture(Path(temporary)) + fixture.plan['transport_compatibility']['installed_transport_versions'] = value + fixture.save_plan() + with self.assertRaisesRegex(ValueError, 'transport versions'): + fixture.configuration() + + def test_wrong_tokenizer_bytes_hash_model_and_escape_rejected(self): + for change in ('size', 'hash', 'model', 'escape'): + with self.subTest(change=change), tempfile.TemporaryDirectory() as temporary: + fixture = Fixture(Path(temporary)) + if change == 'size': + fixture.seal['files']['tokenizer.json']['bytes'] += 1 + elif change == 'hash': + fixture.seal['files']['tokenizer.json']['sha256'] = 'b' * 64 + elif change == 'model': + fixture.seal['model_identity'] = 'b' * 64 + else: + escaped = fixture.root / 'escape.json' + escaped.write_bytes(fixture.tokenizer.read_bytes()) + fixture.tokenizer.unlink() + fixture.tokenizer.symlink_to(escaped) + fixture.save_seal() + with self.assertRaisesRegex(ValueError, 'sealed model'): + fixture.configuration() + + def test_library_and_normalizer_configuration_rejected(self): + with tempfile.TemporaryDirectory() as temporary: + fixture = Fixture(Path(temporary)) + with mock.patch.dict(os.environ, {}, clear=True), self.assertRaisesRegex( + ValueError, 'preparation library' + ): + launch.direct_worker_configuration(fixture.args) + fixture.library.unlink() + with self.assertRaises(FileNotFoundError): + fixture.configuration() + fixture.library.touch() + config = fixture.root / 'normalization.json' + fixture.args.normalization_config = config + for invalid in ('[]', '{"enabled":1}'): + config.write_text(invalid) + with self.assertRaises(ValueError): + fixture.configuration() + + +class DirectLaunchTests(unittest.IsolatedAsyncioTestCase): + async def test_two_children_only_frontend_gets_config_and_requested_affinity(self): + with tempfile.TemporaryDirectory() as temporary: + fixture = Fixture(Path(temporary)) + available = sorted(os.sched_getaffinity(0)) + fixture.args.native_cpus = str(available[0]) + fixture.args.frontend_cpus = ','.join(map(str, available[:2])) + worker_observed = fixture.root / 'worker-observed.json' + frontend_observed = fixture.root / 'frontend-observed.json' + ready = fixture.args.receipts / 'worker/NATIVE_READY.json' + worker_code = f'''import json,os,pathlib,signal,time +signal.signal(signal.SIGTERM,signal.SIG_IGN) +p=pathlib.Path({str(worker_observed)!r}) +p.write_text(json.dumps({{'config':os.getenv('NEMOTRON_DIRECT_WORKER_CONFIG'),'cpus':sorted(os.sched_getaffinity(0))}})) +p=pathlib.Path({str(ready)!r});p.parent.mkdir(parents=True);p.write_text('{{}}') +time.sleep(60) +''' + frontend_code = f'''import json,os,pathlib,socket,time +p=pathlib.Path({str(frontend_observed)!r}) +p.write_text(json.dumps({{'config':os.getenv('NEMOTRON_DIRECT_WORKER_CONFIG'),'cpus':sorted(os.sched_getaffinity(0))}})) +s=socket.socket();s.bind(('127.0.0.1',{fixture.args.http_port}));s.listen(1) +time.sleep(60) +''' + original_spawn = asyncio.create_subprocess_exec + spawned = [] + + async def spawn(*command, **kwargs): + spawned.append((command, kwargs['env'].copy())) + if '/test/frontend' in command: + index = command.index('/test/frontend') + command = (*command[:index], sys.executable, '-c', frontend_code) + return await original_spawn(*command, **kwargs) + + callbacks = {} + loop = asyncio.get_running_loop() + with mock.patch.dict( + os.environ, + NEMOTRON_NATIVE_PREPARATION_LIBRARY=str(fixture.library), + NEMOTRON_DIRECT_WORKER_CONFIG='/do/not/inherit.json', + ), mock.patch.object( + launch, + 'commands', + return_value=([sys.executable, '-c', worker_code], ['/must/not/spawn/bridge']), + ), mock.patch.object( + launch.asyncio, 'create_subprocess_exec', side_effect=spawn + ), mock.patch.object( + loop, + 'add_signal_handler', + side_effect=lambda number, callback: callbacks.setdefault(number, callback), + ): + task = asyncio.create_task(launch.run(fixture.args)) + try: + + async def observe(): + while not (fixture.args.receipts / 'LAUNCH.json').exists(): + if task.done(): + task.result() + await asyncio.sleep(0.01) + + await asyncio.wait_for(observe(), 3) + finally: + callbacks[signal.SIGTERM]() + await asyncio.wait_for(task, 3) + self.assertEqual(len(spawned), 2) + self.assertNotIn('NEMOTRON_DIRECT_WORKER_CONFIG', spawned[0][1]) + config_path = fixture.args.receipts / 'DIRECT_WORKER.json' + self.assertEqual(spawned[1][1]['NEMOTRON_DIRECT_WORKER_CONFIG'], str(config_path)) + self.assertEqual(stat.S_IMODE(config_path.stat().st_mode), 0o600) + self.assertEqual(json.loads(config_path.read_text()), fixture.configuration()) + self.assertIn(str(fixture.args.native_worker_socket), spawned[1][0]) + self.assertNotIn(str(fixture.args.listen), spawned[1][0]) + self.assertFalse((fixture.args.receipts / 'bridge.log').exists()) + self.assertFalse(fixture.args.listen.exists()) + self.assertEqual( + json.loads(worker_observed.read_text()), {'config': None, 'cpus': available[:1]} + ) + self.assertEqual( + json.loads(frontend_observed.read_text()), + {'config': str(config_path), 'cpus': available[:2]}, + ) + report = json.loads((fixture.args.receipts / 'EXIT.json').read_text()) + self.assertEqual(report['transport'], 'direct_worker') + self.assertEqual(report['status'], 'STOPPED') + self.assertEqual(report['exit_codes'], [-signal.SIGKILL, -signal.SIGTERM]) + self.assertEqual(set(report['cpu_affinity']['roles']), {'worker', 'frontend'}) + + async def test_bad_metadata_no_spawn_and_socket_ownership_preserved(self): + for error in ('metadata', 'socket', 'receipts'): + with self.subTest(error=error), tempfile.TemporaryDirectory() as temporary: + fixture = Fixture(Path(temporary)) + if error == 'metadata': + fixture.plan['transport_compatibility']['installed_transport_versions'] = {} + fixture.save_plan() + elif error == 'socket': + fixture.args.native_worker_socket.write_text('owned') + else: + fixture.args.receipts.mkdir() + (fixture.args.receipts / 'owned').write_text('owned') + with mock.patch.dict( + os.environ, NEMOTRON_NATIVE_PREPARATION_LIBRARY=str(fixture.library) + ), mock.patch.object( + launch.asyncio, 'create_subprocess_exec' + ) as spawn, self.assertRaises( + ValueError + ): + await launch.run(fixture.args) + spawn.assert_not_called() + if error == 'socket': + self.assertEqual(fixture.args.native_worker_socket.read_text(), 'owned') + + +if __name__ == "__main__": + unittest.main() diff --git a/src/tts/magpietts/CMakeLists.txt b/src/tts/magpietts/CMakeLists.txt index 21b6329..e44b189 100644 --- a/src/tts/magpietts/CMakeLists.txt +++ b/src/tts/magpietts/CMakeLists.txt @@ -14,12 +14,14 @@ add_library(nemo_speech_tts SHARED ../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++ diff --git a/src/tts/nemotron/native_runtime.cpp b/src/tts/nemotron/native_runtime.cpp index 1fc2526..4956e09 100644 --- a/src/tts/nemotron/native_runtime.cpp +++ b/src/tts/nemotron/native_runtime.cpp @@ -6,11 +6,14 @@ #include #include #include +#include #include #include #include #include #include +#include +#include #include #include #include @@ -27,6 +30,8 @@ #include #include + +#include "native_worker_protocol.h" #endif namespace nemo_speech::tts::nemotron { @@ -180,6 +185,8 @@ class TextPreparation { #ifndef _WIN32 struct NativeRuntime::Impl { + using Message = worker_protocol::Message; + struct Request { std::mutex mutex; std::condition_variable changed; @@ -200,6 +207,8 @@ struct NativeRuntime::Impl { 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)) @@ -226,35 +235,46 @@ struct NativeRuntime::Impl { if (::getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &size) < 0 || error != 0) throw std::runtime_error("cannot connect to native worker: " + socket_path); } - 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); + 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 (...) { @@ -272,6 +292,121 @@ struct NativeRuntime::Impl { ::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; @@ -321,7 +456,10 @@ struct NativeRuntime::Impl { } void send(const Value& value, Clock::time_point deadline) { - const std::string payload = "J" + value.dump(); + 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'); @@ -365,14 +503,24 @@ struct NativeRuntime::Impl { std::shared_ptr lookup(uint64_t id) { std::lock_guard lock(pending_mutex); const auto found = pending.find(id); - if (found == pending.end()) + 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) { @@ -451,12 +599,95 @@ struct NativeRuntime::Impl { 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 (frame.front() == 'A' || frame.front() == 'F') + 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))); @@ -479,7 +710,12 @@ struct NativeRuntime::Impl { return; } try { - send(control("abort", id), Clock::now() + std::chrono::seconds(timeout_seconds)); + 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()); @@ -508,6 +744,10 @@ struct NativeRuntime::Impl { 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); @@ -521,7 +761,26 @@ struct NativeRuntime::Impl { if (prepared.apply_normalize >= 0) submit["apply_normalize"] = prepared.apply_normalize != 0; try { - send(submit, deadline); + 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()); @@ -591,6 +850,12 @@ struct NativeRuntime::Impl { 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; 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/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index 2abd29a..14e933b 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -138,6 +138,16 @@ 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_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; + } +}

d7TnOS{yo!QOPNaVqF=Pb^6Fi!TFK$v~kgWE@AF)si69h!)B{m&XU z73=!fx|@k{+V?%!8xvp)UrA8sk~;JP`><9`=KD^a^$ zM|W#)E!OqF8=?)VFU_VjY;B2QKWXQg`&zucXO;^&I|%DYEb(Vc7~hE)a&{&*#5xUI z7hzq6-9`-kbraTISPx-6iD3+JVZDf<&fdcM2AB;KH!@%qq~VGJ`#9PBP(cN4>W z-a`yI@Ab?*_dLrpx6awZ?h`gg*j!=rgxyaJV}C%{gTyc|^MySm>|tRGggqi`A+f+V zM%dPS-(6zM0=o&@HewjxB9aqx_9!vLEhZKudW*%*WV0Xme!YamA?IVnP|H$ch+8IX zxvCPKg=aQP7{gP<(BDd8h*<4PhIJVO};7!@O)J zhPW-l)cNJDUYvUlWSg+_!uzkWl6m~_v`oO+W8}Qf3n|@2j3N%#TG^lX6J`8i^F$MvD*A|MW2Xq z_~(N@7307;9#$*P+c0CFcfc6&c?UKQ|GdlrQX7ts2Zen`EQsd%5HaNZoLFK}82be= z#2qGv{yH_cjx&xe_8a29Byq5>h$R+<@kfLmC5D{Gh#}{3Vu<@%*f+wy74{vm#G-I~ zCxm@Z40Zk>>_=fg3Hw>tFT#Emc2d}H!hRR_hp!d zu^a@i9s&Hhu?~sDajCAbdcx`p zYapzlutvl%zWm~xTVoQ3@ih_FR9G`%&4slP)>2q2VXcL=5!RMiV7p1!@|Lq(zQ431 zahT8c!a4}+D6EsP&ceC~>niLvVcmpv7uG{qPhoLky@d4^)<;-hVYdtGM=Y>iC~WUVIzf&5;j`c9m2*48!K#_u<^ns2%9Ku zlCV38;W=P3F>LoK!lnwFCTzN}8Ny}?yGz*J!tN1vFEPy7EMl0m*~0DZ8Ve^FD zFYEzf4+@(v>>*(f3tJ%U5n&63Eh2{P_b4%Jzs1A?I~azoR0q4|&+kh}9CAJ;Y^kti z!j=;Yth}(5sbaUh&c{g{>I{TEA?!(ED}+5IY^AWLiDBDzAfJ80_?{tguvNlV3tK}B zwJhzO`F41g#09o`*q$SXI-e&NB%)%m7lYvL<6!uzX99PAom zfdz-{v&3z3PSkQOiGy7yte~*#g{2BBL=5A*L0Dm87{iUiiV#D8M_SmV@i{9>;?Q3) zV#t}NP3Ado5^=?ep}!Ku&|gVmwaNeYLfp+HF0j(WR!UfLk+U?3!x+jCLoH>6wdr9M z@oiL&#G$|Po+S|k@I;GS6{?65VlU_T<@9t+M}V!*+^JpVNHZJ6}H=} z)7_VviMZy%4tP16rth|i+ZH13uovgn(o)2=64qK+8)0pQwG;N87(;sz*FjiEVV#6^ z_RKZ_j@KmoLz<8CajyVOAFYD_!zp2xE{iK3X2QNFKX!( z#idxeQNwny|3Q(8h2ec~v9AC9C4IblEDiGF?+v^TcMYEfVP=1e_^-2()BoF}uc+1k z+v9ex)>P{zYW+kW`Rm^}mB;pjg-c-yDO)y8dTS28(t5&!7wu z>-yhGGE}VVfA7dJv9AAj*Ko0}|9931v9A9;BqPPT{@-MyqIKQEWbt30%5(m3a z*c@SVh0POozpw{{Jt%BGG0e+D!X74ud08Or5n&67p_WC&0^4k1yR4_(^6mF1i9^oC z!j=%j_#P9sRM;|Nft4P%<-#5p_FW>Hk3Ar97{e39FfUIML);2sPYGKo>}g^_VqD>8 z37gmte0X}s**iG18BZe`k|6jD0#9<7t3tLAF{jC?aff)LGL)b=Pn~0$n_FY@B z%_44#XYTjqZuQLOA7j`?4E=2vHd;I$(nuWo+ac^tVQ&#bojZjsX=A~s6JkLk zD*WFH-fM4Lj5?ofZog%0ch`=$IN#9Ur=&mF0bvJ)eJ1RXu+NEMd|wblorj4b?n`2b z`%2gmVMm1>6LwtK*TTLL_N}n*gqVOBM_6k%C~WfPX27{*>${2!SdA}*(}T*59OhIz?N z40Z19c>bKhE+lctc@Z(h^=Wl}PQ+a-;$HV&zgQN;T_WNx6?Pdh>1-7}v_LBJCf<`0`wKOJ%{+f8^9($U4W;GzL znXu-bxj9>S=H_fk3}a|T47Ict)<#%cV(71Y;=~YtjI|?isHMHI4#GMTLoJ<%p})?= z5O;kGn_WJJE?!(+TQzE3J#)9;ZJt>O`s*gFJ28x*hp?W+&|jPwa`qC|n;3HbCtjoU zA#v!hFEQl2ofzW!3F|LxfUtqYP|F};gFSP%`w-9E?LL9LuN3_aC2<)0Fk*-sE^Gub z&haD;W1rwz9ZMRuiJsYF#7!cG{^p8p zdnbuQ&dJ1(a|$uUO%*mx*mPp(Zw9d-nwLU}FCtGqMj~z|i9?-t5kuVF!tNo4oc9Ww zMeP4$AlGkE%WM*d{_Z1&xH-b+5<|{;!tNLL05R0^ps@ME9`ek6{5|Yh>1cnjy9MGs zZwo}+Bf>U`oC`(VUJ?KWY{ zMb5{C1;S>Eu|FZ=o)osiGaDKn5B7dZckHg=f2+CgF+&{tZfcmlzb~_#_I^IsQtj8M zJr&iNf(79{a{BMVTPeoGzEc``@SW0ZO!gi-w>J9?$D*f6f3RnStrE6c*cxKPt)Q?y zOANI^)*nsp>VlNRxEiVguMcAvv(4Tsr!fPUKRd*{m<1uc(q0Y4=4)(gR zb;8yY!x+kX=cV1}k_}#*`+1%>Jafme(K9#aCSjXBb8~JXhB0gvwv8C(bGxuK&)hNW z@XW3AO<`{d+bL|9u(yd}Ufv;wvA-+qJz?(?Lw~!4?Gd(@80!2$*gj&Y6P-0gnaGdJ#xXKvhCVi?~! zVi?~)#1Qu{F~t2RES7lj$M=^k#L(Zxi7x=)aVaDYwPY2RO;~nfsAWRu%=h^mBo4LY zB!-;1h#~F*VY!JRXBBeJ3^^|paqRpS;x6*yO50RK?PKvb=EWlJ5@DAL%XPJNoUu>1 zwOl6RE*F+ZST)gKUJ-YNuq%b-6P91t08wWFFV1EV`wRZpdWtQK8qEJZ;#DFK`^+Ko zu)oD&_Ble^o~(cWJ4>#P`e*CdD@BN@(m;|Evm1uup5OHA(puGALcD8tQfJx zqA>0zVyLsYuoA*b3cFcYDPg6F1<`!2CfC7Xf7Gy*GLDn>TjHxx!}~HMC&pe@SUJz! z*CORTv;G3R30nnX*nSm>p_WR-f<(qx>~ItNA$fd=t4!ixRfJU~hB~VeL(W@-RTox+ z82YP84E@y-b}KRDtW69#>jTE*{wX_x1ju>*b_pF&sOw>AfW{c5ZM`4|ap})?;x(Mqk>^5Sk zvm3D>8eex}sIvz#xm+6lCV38 zVO}N^LoHK?A#SR$X`Z>yW79ozxBCoXGlktn3}ayLcY@t5;_e}avEM6fmay5v?jwde z=Lqx9iMYqsxg-u_m?!Lh&)oAD4+uLd{=R=u#LX90;F{#egImju!X6Si9~M?gt zUDROy`KOgCPc94QdHnC+eLBj+`t-j8>lv}G{~cJXM15<-IQ{RyS}pR}_geYSVVE!W z-W3?$yTZoopKn?t#_XSOdRDCKe^1qOVqO0|GS7>3{c}|>h;{w%sd`bY>wizxOVPUS z`)zQHe_7Zo#75a>3)?L3eU?~m_9}@(f3Fb>EGTS2PP^r=Yu1uD6!I|xZj8cwnEstR!S} zmyhnDQv$OUtGlX64qN-A7Ook-7c)3u>QgZM3!Q^M%3{5^RM3w z6zkgmGf72)@P3fjrxEh9|6v02ugMG+efsY&9wPdj5Y}h=J{zZh{bs1hgX=dekN;ld zVWM_>uW_oC6t&@EUH`quBSd|6%_g&P`qy1XiahwNG8?CVon;islNOC@G%>u7^$ub6 z#e1#Bh`6!DP|G-CsAasc3B-_dqGv;+aogP_&ukgu?j(ju$G!sZFPU)Tf0 zFus0q+e7&N@*s)B_~r|HNZ7-|76^NUSYV}xZCbI!?NoLhccF+|B_IBW#stb)!1$ZnbB&3^~^bdsf(U!k+gG$7<{M1<%}CUKI8cG0fS^ zp1J+K;+fmutHhA=HDTw-c>=_(6>;hs@9SRN9ad=6)_LacXX`z)5Y)24Gk0Fz@GS9> zmZ-mtp4npTFPn&=zsp}*b2_7KDP_7X$R4}|R_hMXS~L(cug z5ciR=kBK4YC&Vyk2U^?YWxPg7T;_@9^HUOsTAn1=FVNos5(hg-3^_j|hPXq*J{R_d zu*1Yq%a_DZ%U8mV2s#90IMk9`*oDF_B8FNn7IulSONpV*%YU*ux`g*LY@&(ciVgt`k;}80x&982U@~EO81ls@C<87rLsOf`&JV=t#IMx0&0bGM&scD*h$ z!}U4z=f>f?fKqH>)X<-Q?XF2w8(Y`EcGpyl1J~|YJ$CIbvvJsW3YoF~?K)nnanzcL zTK(@EYA)9GuK~6Y>-yLAT6%MkisHiiRw9r8nXcAiUH@}hZN$3%^}@DZeO4Z}J-c2A z^RE}SBY6T#7Pj`pa4ha1tfR0_!a57PifC-bM^@-G~iOA9}1iF~s#ChW>gA zixcY;_1B9S`s+;${q^yzWt7kE`g)dF7WwVOkh7n#{=`t{0Ak2FP}m@0gM|$thFXRa zLoLIIA#OM^#ElR(QrIYAqdjxa*WBTmZ93F4M%Y+k$-JV%a)OnAv zdxgy+hB{{xLx1-XL)@2%FD2x2J%_|0=Uid)gxxRf0bvgcn=kAkVGj#iAnXxg3yEQ_ z7ZJmJKI)nKKiC(0*1&eAs4el#79;0l#8AtR2sgzXpa z1ACUlp}*%mbLZuG&)o67AnZk9FA>A|UKaL>XYRGHS3PqdN3VJ2j$y5^*M+SkhB2%s zhPmD#>7(<${>fU41{eK!eNF4fm zlUQH{h3ze3sAVTHSJ->P-Y16swu$q|yGdM-n3Y&;kFcfUbAx+H z9LD#7uzkWl6t-X3N5Vc9_KC1hg&h!fP}pb0u>B4Z!+d^D3~^r&L)>9uUkdw*82an$ zJvMDO!oC!%X}@JWr`m6*s!ThV3WMY#I9dS=cYcFos`=A?Hb9zX|(Y*dN4DOBwRMNYrx5 zi*wJR|0!&Tcdpot+b8UEVSkDK{uXvx*coAGg`E@jkFd_scZlIUu{$sSin#xT#YWj4 znQ<&|b7m2iBCLm~v!}4EBIoq$Gp{q7h|4Z4hp?Q&9u>7bC2XPCKXZxxF7V7|0*^oc z+KFrI`Uv814J5?^q6YK7V=lMo->!+c`;YyGJnXs#j9u4&vFjQz|GLJ7QEhA-{&kIu z#Jc`}~UKL5MwE)jX`x`xditKI)@x=TeK_J0#lk6pJ&wQizznV19rI>zOq z&u-q|t!>dkc%O&F;W;6%uq%k+T>h2BkTajK{KSy605RmeO4!xJkn#9c=W z*G&o%!{@NB_ssp9FV!>KM5v_@G1PJcF~k)XcB8N&#L!<+V(70JF~r?O3~|MY1vUX; zJ6XeSGoJJ8H^h}7aptj~L?A|8}h};?(ya zHXw1RrJ=A!!Wt997}k>W958ilqzTE1I-3$h&St`z6GP4x#E`QkF~qeZhPc+m5Z6Xn zTVlu=#BDF<$H8_a4*j(k*1Soe?2{O=Pd4-%^7OxC9F3w^w-C;c9t}1eLb_q=nrUNk3<0dqFUi*vWpV9#tO5I02FP+`M}VLpcwLw_TP1y)el zMhY7xY&0?ScZaYs!o~_4Cu}@1%*zB}6NzCAlZ4$#EJ(~|EH;@Ka!w(JI;Rpt+%#f{ zn@$XIGlb0)b{8@9cek({iIceev2qWIL!I{un&?wy_U+xsr(o%_do`{py}e81m&XXZ`4S&!K~ zk;_Vtxy#av9&>Ad$z!hPD#9>EFB67bRtvU9u(g7{BG@{@kmsv}AkYI-eJ3<)p+)wsRk>^pN`%bX$J%-gi z%j_SC4E1I4bUk8?eM|9}od#W2!jNY+!l1iAu6D}PydJf%qW?yTg z^7Xbp(IJB3NU>kjqeUZ+H`;LoQ7PyOc1*yNob+HWTb}!I~3> zcvldHcr6Hnu5Iu0#~rLC(ZRD7VbGlu=fPSFT^qs>uPtGS*N!mg+7kv{2g0E1C|D=L z;Mtilcy{qv^wD;|XLr@?bpPx~ox-3038tm!p z{O}w?7;+g(7<7%rHRxeP2hZUi>ur;X+X#=@$jkTqYF#q7XRdCV&`lR? zhG5gheepAeZkEu^61p1%yHT)E^vi6an=SHuLwL>+p6?3XO+q(Uux-NguwW;}egCHf zJ0m>PM4mSb_OFOHPv~wDY`$Q(3bsJ7+XQP=J}`~1f47Nc>2~3HhtS<2bax7Nmtc1b zwotG|g54w7V!@UOcCTRf33k8e&-;b$ftaP(A&wii*Z%+Hc~Hzd*~CwUfq8$(i)j&j z@q3pn&ir8PnO|W3Grt}dG3+zH+-25&f)9QR!)o^bKhGngX8YTiRBIf!N5#DUZ)KK> zdHvtYJSOJ#e=GBNJTL28`+tCJpRl_9=Yu^VV*1Yjds58n|8^!6^ZL&KdrH*8o&kpO z^?y6_H1SD~$Ltxwo)zplk3}D2&t)F7lM(NE!Il#i#-59!LCf3*yPhu)-OyOKLa+~_ z3sT&3CDEaV7YS<>#~VJ_{=q$8BD&DdZrUosP{YfFA(z#Htr2W3VTkt%VTiX*uva}c zI-|7c^)-*#iD5*7V1D!uA`<%yxhx-JXG&zqWj*l(dtI;%gdyHW!r=LaU~duz&$kGJ z=i7q4BiOrwy+;^wIVP6*bfJ5nFx0R~unz>=EZB#FZ4qp%U>^zgv0$GF_NidoJm&r% zl;M-;#J>2L(GM*kQqr2zHb( z)bMGvgOa!`fqf@*4aBvG?}hFMk3}DC7n^p>V|E&HIZhb*^8{g`6`PjY*uEy#7urdt z)e-l%{YX3!?MC-{Vv!ag8eDjUxXpgzX?O0{|I(cuv3Da7VKZb zkY^Qg{RZp|(V>R_1UpL@JpKC+?D7TAb3_LgME|5KaruO27Q&!Q5iBcV@Z5Ek%`7n< zJhKViVsh;Rp6a^f1;i8avJ(c+93G3h(57x$PLJ7X(9IHiOr_9CW#{h))0 ziR%Y!`LWjz{4ojbB-1(y&rXD;XUKxigoQC%+Rpxw`5kCKLDz-oz`6?7O|b5Q^$@J5 zV7&+%Vhb^?H({uukH_5Yd0&rNFX*lmte;^02}7O(1RF>ga=D5ycxI1o>gHp6HPONI z8o{n54DkkeEc$3W&a}ZEv(pf7h+snrL%d;v4JQooMhG@iuu+1I7Ho`QV+9*W7-}C+ z80woq7<3Z}gKm;wlL>?86vE(nonTW1yWV4??E(|GX&$qa5pTL+GX$F{*etM{5F;A0-M%TQ=vOnaO#MzCiILoUw| z2G3=JJx>@s{pSk0EOS9;7&L@qB1_7Y);w@R>=1zSxRa#P>{Y>DBMf=2Ck(l~?lE`&dV|N@ao8x>8-l$_7;<@wFx2q2VDAXlBHAfQ z+`oqByF&M#VCjURhW7>AB-jUnZ6*vgd`K8_*+Lj}TLt@wFnE3}*e8O0D%du`whOjH zu$_W^Mi}~K7h$OVbHR2K2G2c&!E>+2-1YAZkGbpLMDKc*tNT*uz7p(f!Ls+U(KFYh zmiA-Z_6glL9*e#XW8Hp_^@gE&Kj1No(4XH5cF<$)b(=#Tvz{1-!ya?%JK`}nm!pKi z^E<)57wiYYjtO>LuoHs)NErJ1C&7Lu47LA4SZF)Ww7t=fg))~VI}P!EB|3QiCfM(U z!SfHo;Q1$Ep-srNzX*fwZ^97oAHtwJDcC8Gx%*kCJ?5@Q|9UL?8F)qs!5NR)$;jnD zkGb2BvmUcL)OXHfEiFWwX8%YWTg1!aG50rPDIT*fh?i9`^~}L+M2B21APn*RX9xDN z$;2%?(ZMr^U^xjxyj&i0$04`JY#Pvg;r-sT(IGGWkF5v;1mqD3TL&Z>FL zPKIZ7!jMZ1!DJo?!I}L%fR#L%arpH59CoV3!Ei zSgLSZj~Do^3ql zZkO76%<2%YonY+=Lx1)q_n3pR=jS0_2jU6VQLs*ebr!6PU|k7A4c%gvVmm8tVBHOK z_bYn{pPqvC@|b&^)!Sn>1@u!N!jM~E!LB3>p8W`eXMe(=8$cLz|3w!y`1k!lp<~Zu zgD1NW8o6BMdAirquJ)MiPk3GgziR>cZ*;r$?u#HZ@$poD%b+SZu3}FQal#7d(0h&@#HyK=+8TZZp1|E z!oNH26uP?vyPGiNve0Ah_xmD`+03B3hcJu-``rWBVxe0?7~$gCs?{*?+dm`unz=Fsg!*DHVfT{f^8A3u!z@NjL}x% z`H^7Tawe~@gUID$;rWSRp9;23u~XP1wdCaDWeWKaqH#=bd@3?-BeNs@Qg59?bRmRm8Nv4|C`JO~ho+kwcyKS#fq>h2`V_2JCl{kN<4AKSVz4H(|)f|4rDR z!iW7P3_kvE!u}Gu`@aeMJDxYI?V7m3*Z)n}Ni(l|tp1P47rzn90#ox2U;7&|cc0OI z0<*sfOEyOKH(;q2$L*A;(|@MiX_0H5;npy7AJ%?CKK}|`?{VkrLOaQ{Gen2|_x}hR zpCJp*5*FR2Z}>UEg6Y;Ev4-ffikX^)Fw~bq7~*9mY?uvWS~kL>Q}l*iKp5g>Ck%P! z5G*HQ@XSS6bW4oMB{yNHFArf8ZE>buNLbYUhUFy;x_pGCT0_(F6E-RX1qBF04Fw4s z8S4rWc4^EC6ILN+MF>lcSy93;M#Tt2?ZpX0pOqkNQXH=&VaTNvVaTO4Vel+N7?$0# zgoSY~7YSCUq@5{o9g8*{W-9v)DwZ?GLThDOIT0^au=0ZW&oH)eX4x0h7Lsi*&d0Fd zVNb}A1?+dvVAy77iRZO%L*u{g_l@6YVazI!8nBJ5NEmceiduujaR93%bhAsGufzJM z#;7v!#J0!oJ0;7d3eh3Ys)S*T*f>B}jp#6bvq&!J&+0@6&l-e5SCcU4Y6(`Gu*pG2 zv(Fw*)@OBy4%?8rgoPQFLGK*NJ?jZweX>rY_WDEzb}?Zi<6PMIfwBEl)X;!|&w2tzJy z1^X)6$R=*D;n|MpFh(p-#A{DyiUypDw7J3t+?sO9C6Jf>5Vd>@q|+a=`EiR6MB zIt$i?Ff5=&?Q9>+Os-9(<m!2*Z9ITPCm{$?Ah=Z{j({WNsVl>o6ig(1++|#EkVB#$A1{9248Rz9b&T zoz;%-aJJ8f8m=Us*w(S{e)K1+598O5c!KpO47m&-47!1YjgI?T_2*SYhZ@w_UQKk! z^BTdf6>N}Tg9RHx7}i_0&pwps(4WHy!@dUFPsVbl_CtmfPxRRc!fV8@q$et z4CBsnf#*b`!}l~BTkKye&q>4+=L=RvC-oD*SJ2n&`wxum)8e~BE!Ve*z>jj%8So;EjX^DM?c+-V$24OgdG?OrR&Jyef!EO|6HeuKfs`X(A*{_TK zHL~du_C1E>lPwciUf4Q~ai2r-9AaYG#@2^0B0=z_Sk7)Dp3pUoHbRMGi(KXkT^eD~ z-Aowb%_9u;-9i|2^98$=Fyyj;FnE4K_LuNIrnWD)3D4UJLoRm^hFtCx>@LCXh!(iS z@q_2xLN|~s*U&8_I-Ij&`*rwUVaq4J1MqxF_qz)57Lj;h_Yj8lmK{$+S2x;lC-w_; zY`ukUF^LDZM6i1WTSfXB%fx*`=ieg~n1*$QZ9|af{lar%CF_#dFUaKqqC-5@XAcq` zJRc$qc|I)IBZ56j7}i^Rt<|p&@s<)D*kghfb$2lRI)6y%L>NNUQ<-wCX&{4C0CH#8V=drv-b4u$i_SO;gM6 zvqXnw_c_AC$T0|(5r$=uZR_CqJkh~(xnM5{wt_J1Ypf&;@m>_{CBaq^hJ7OT9e{W* z6CIY%QSEHUCN4{0>Ns^Z@kG2egyFl&-Us5xgRbm_=QE70s&6fchx2Z%1}yWqcf|c} z#PaeA@kE}3$oC4^I-z@&FvNR}FnF#f49Ahq66DV7(Cw~ z47xW3dy6o5zD*cB-ysaTcLjS-uyn!RCk*EYHxY(BKM-s)VWZ5**v9rXqTNoz*tRaT zBGYcqV_y@;HnfvW`;f#7BN7B#1luZDn|#(FF&@@iwtRxw`>K3hXeXJri+F|+34)JE zF6irz34`tv!8(#-ZEPp(J$Zhf*gyMJcy1GHyI?OBN*-?q(IJ$CpwJnZo*JsJ#p@EkI?NU4Bs7J5QaYclCT?W7n}AKVPQmq z;A_HgZe|u42k7<@9gYpzJ|Ft*8=`~fe!&h92G4H^gXclPrjz|v=r*JVrX_Ae5bqH2 zL|@<4I8o<5ljN|cYh`DR+Yyi1$zjZn3icgg$n$%`;JJdVe_%fd-7&$A6NX$)2=*gk zSeAYg>}SGoJjRwmtSfA}M!a8$C;FucIhTO#lKTF8S8QK?74bHSYoEUf-IB}ff|I!H zqK4mz4rBBOVbJ|4*k6MEE!aPV4YjG8#?HZq5eb4373?30^xkhG6X43)X3NJ`#CmCGo(r309u0@6cU9bU5~9d14&06Wz!lqp2El5FK(^ zLY~6`mQ&~&q}sTA{Bj9hZo%>h)|{-j$g@GTFec_Q)TU<|+h@md4!gfKj8x{HZjnn~ z!SV@~U$6p#6%;Ha_oVnk}tfpYK2ty6*SQL5ICOWV>ghi7~pVcJ{x_W}u7wnuEzl(`3nq+cm zKp1jqC|D!GE+Gutg2sd)UK7EZ3U;YrmkHKPu*(H&F4z@<%@fDdErhP6V66meEm#}D zW=AJ;6Tbs6er<)WonVFYn1#*zUWsNmQ`-w&2f;cD)`>9G&{?o9gkkyYDp)tcx)T=d zR-4+{K3_DuVLgZrxhy9A3HAhG@a!o(dkNNCus(u)OmabWY zkmo?5TNz!PN?ccU6l|i1w_kiWP7=Dw zg6$N`#1x{7CYd$pI>NB*P9+Sw>j{Hynqbofn<3at!lFs0hFOFmmm388PK?8iLN{Bm zIfC6J*j&N7=Ql+r+18~A-OYl{6YLhj<`af}javm+<7i1H+^U>_4M{XCo zGNRA!5V|`ByGyVgBu~_LH_@Tb77DhAuxOI0{T{*)Z?RxY1iM$T`vkjRurI_oJRo!r z3U-Qk;<_O_4o6=a(PGS-WBXU!O{hLU$9Mr%@;L%Aat7r`%thgf^8M- zBf&lv>=VI073_1;FVXB~ao;9%+XdSp*iOMd6Kt1Yp9{8Iuswq973>Saz9bCa(_ax5 z+Rch-UkkQRuy$jtI&ptKw3AHxM(FknHj3P{hZ+tD-Dsi=0qd z^Dtq^<%nQM1^Z60?*;pTuxOI$n`49_&*Orf5bQ_6ej+TIWODhLFyzu%jNdOpSE;Zu zPO^OdN^}^%-vs-eFy#4%V1EksmtcPj_K#pE1?x%r1+||Ny3>OFE7%#q{uAu1VCMw8 zo79If3Ze_7e4CYpuqcdKic$!JXI8C*#*lXSWdxm36@*1Jc8{MeSM+O}@9|*h!|bd)C8<1VJI;Sy-?lf^`=6w-yz;s-yN6Q0Ed zDD|8nLR!*?(B3`P{l^3jnU=;;xB*wjx&{ZZ3-vLz!!?ILW zuxcKg8m}kzt-8nTG{mbxSZHT3t)^hL1gkAr9l`1fR*x`zKh!5Ij4B9%iwT3S0b$TJ z6zm-_?xTv@z=_{+=<7zp^Af=t3)V!irh;86*kyt>6YO%qnhSP?U@Zv4^TE__uUZlv z#`YW;Be1Dte;3P3E8*FiFzDJ4hInlSTT;v{O5U;rT|1#`FIWe`IueFuw-aH=;4s1sfvRP{MG(KD)mhQmd++eI!bh$m|3IXQXlcL>ir1-nbIoMn>7D?k{wJ$DmN zu!VvxA`G?PBiLfWmI!vQVD|}jzhDm#hT0z_4Bul95r*0yCJeeq2!rlX!Ild4Dj8dN zJ|=XJ3-*LyyG6Vwh3=qOS3;u0I6Ord>U&zS{N#C7&^;q`&niZqiv!OUf;}fZ+loAw z3Efr^Z?9nc1nVqvd0ynQT(B#JXEE_yj{ZXT0*QxbYpG{I2iFCr@$aJ*#1mt?Qm_{V zdr7cWg1szQ;nd_ctQNX8f>jirl?7W%Jh82Ng|N`3Zd#L6`^v|mnP5}IGOucAv*NuS3>u-VEY97MzGgOE*QW4L^m`ZTlOp( z=nfDa*tdcm6zq^-hXwmx)OSSajtcf2Vb~^rPZ(*=`zdj|3%0RZa`t@Xov9=C7<1ttFpU12W_3an654pt7 z$;W+h-Q=u&&E)JC5wC;LT~H&r=L2FnyGZEHkmq2bU-C3cjMvwGjN4h!pQlAG>%4Kd zc8HheF?X3bC*o!N<5oUK!@OnJ^~@4IST6H_oxAHwiec{Z^GZjXX=V(UWfeXb80Kzc zvU&5mJ~<3?`#5_%Zx$Qcw5}O%(Q8hRxjwl(*2&T3j#)}P&DCu+HTH4l&ExsBa@d6) zb92k<#c*}oO^gnX&$hT+Q_@Sv)WL!=f1Kc`PDW zd(SgE6>Cd2t%Jwxv~l*uH1)jQ$BNrI68jAIW2x&*2TPo<3+*J+elKZX6Fo=SNv5gi zXr3bXrl3DNlKQ|p5f(<#%Vf#ng`mAeVu8Bs%b zqC=kQe;4#1I^@}foCCnMNA_$F=$49W3Oz|Yj9)Lp;MtonZ0GwB2G72P!E;1vU|Qlb zG0aXfP1ScL@q|v@Pu{+~bxw>I+DWGMBc5SIf}p=(0|-OE{6>z0!LIaAE?dS9{Dop1Q_k{T$u39&?Y826@cNFz$l|8$uY4Nrn;z&tZgx zajwG!8zESC@%+D$LN`jV(S)IfF@&Lpv4V{gY&>CzH-Rw3n<&^M!r(cXFnA6Z`^n`6 zn?gJh?>fQ|?<3*~-Bh8wo-pX95r%lv1)Cw*W#m~6@a!zuOyN08u!Um+(=xXc_ETtI zO#8LIeNAK&?IhFKacUTmAh?0VgH9b&-bi%toGsWK!D>W1Rf#pgbGZ1PZY%apZxZq5 z3YI3=&4SGn>=wZWh<(%fLU*fR3kbvdaGPMa6NYiPgD`mBNf>nZlnhMcWBYWpz;Sk$ z@Vr~F?%w|H2pc|b3x#fxVD}J)^C61~3*&f81iP0o^yhtqA>RFjLH7V*$mK!79un+f z!jQ`&gdyIegdxwRgyDP^+vfw5$6Uk{xjar7a(RL<)W^1i&^<|Xh&Pv<2ZN_Ne#n_G zFpaMdWdsXJF39sK!jR`s;rX=CJtNq&f<5Ok_gH(G$81lahUW=G4eyV#bMii0E_7Fu zeGQD!3q%LDf-vN=k}&9ABn-Nj1lvODL%*yNx{rts@m?l6c&-*~jbLj9dxbFMvW_t1 z@+x8I&({cp=X%1RdtI;%gu!#8U~dTareK@MGKkvWB08|Q2}7Rm5C+}5grQ&FBMhGD zgh7|Lb6^@@{|<=y-WQ&m2ty4Y2)3Cp)bOETTL?qEt%RZWj|fA&j|KZguulm?F53j# zP8f38K^SV-Nf>mW3AT$ccz!O}Zo*LC9>Mks_Jv?y5{5j#A`G>EE!aN7P{TKZ?I#Sm z93Tvy-x3DhLBgOrB-mk(^$#-J5sz6C#`q{C>&)nkJ@H_JT zhj_<{C)f$X5bsCAp!LnvLHD;{ zKZvNhOXE#zI|>yHI%MB@B7yBMkNBCk(m*grUBIf)x_1uwa*qbCpGe zt|(!sp%`JPA(e~+SaG2%K^Wqd6s#0sh*z30c$Ogyy0U_GB>f4`i-fL$_--smbnr|i z49i)0!X^b7ja}DKoHw*~O5L3HTP znu66L47qeG9GJ%A)h0T`t3wz(JCYi}hLU(-b%p0dqC+nAhz@;LpD@I`*kkVfq76J2 zm=f;AHC5|6q2XN^7P#%m&2Q^HW+rGi~X80u^0F*nc4J?7@woG|1PlJP^lD~JxP z1!3@P>9HO*nYgv`n4OH}r8Qw;tZPFUbZrSkp6v*Ou03JUbr7s0Vesrk7;5ND7~*vi ztgFYO(=lpYFQA8ezz@DY=gm>}I0FayE}J=x!kl_bssP3v_MBGkKt6*O^iKd=d|y?-5V16NI6@ zTZQK!(r3tJfzY)}4NT+P&)bCVcEXV79fYC2I|Z8&eL3-5hLJo`-(AAi15M z5FO$@N*FxZc}(c)ML&Gv%h@`z9zpkOOPg<^4r8>GVbDDxSR?Yge0V-7 zbbW~q_Y$-7OUNY@o_7&X7qq8=F3I= z`Nq?v@c-5#eog_))jr?A^?kvMS?B*QL9oJOZ5_7KV-}&tH)`5Xe5@Kri6$~EH!l*; z&{~=Hl3=R{gXhbHMU#x@bJ?wPVlMby}Ab3--w*-5e zF#MKSeV^T#%LY!Y4?6apg&N)=@leCNg1skLI$^jERQ+DDQ10_ z!|nrxZj;cd>wq5+9oEgwgrToLBn-K%$#Z_~U|Wa|@wO5M-Dkq{Bcel|9}@=8PY8qV zQ^KIzMi_M41=~RwJa-ZX&(8=$pY0+Hy3YmMO&D_7BiKXao?^t?D|9QxJ-%NE-Is!W zCD_-3Z5|bvmbiUE?fZxh^?gGabo&MSm#pvbJRo#OMV{XZ-9f>Yj80zr0rLHa`VI-t z!-@&dk~xy+Qd+Px!ZRDWPXsj_A$j7S9CiNtDAB?5anhggWc%P4_wR@&?jd34Yr)hr z6}~5)$aAykmmh@gm|({RdzIV|f!a?H9oUb8{X`h*`cf*bz~CKe0cbRd{9-?0ez4zWDjI<605hzXntNa)HPtyI?s8 zi#DkyMYi7x#;y@!9C8v*jC(G@auWv6Ql-xCGq5~DccEZ;2}6B@$vy+(%@MieBc9OZ z7p#C_1qCZXu6rO}A)zZQSS{j-@hc*9MFkr+IWUcHdx{BNaluLuHr$3WjUB6DdsdR@ z!blYa$H}z~u*TvU9HodS`m8iz@LVQ5%LrXr!VvEw!Vs^VV5x+`vpiw&tRUFd)W9@8 z?iGm+@hTC9c=N?In94*4&nkq$vnpZGRTHc_VeqUW*utEqTaxrkO`-d^d~#i>3dvc` zXyNDMP)o$CEm$4F8j5&zg>J8?p&rp;onBBaqxKa07B{e)&AH$%m}uVX3!m{-?2MWH z)WgwT?CDxLtbxZYLi|Uflg*i)F3aP^bJ>C`?59j!H~TSe4Mn^Xo~Ns8cH=d`<{;q02W#;KJ_FPqXUQs(y=k{mgI2Vkk%bFVI=H0~0+hvz}%=NjaOsZd!K3xkGZ;*92nJ~obB3M_#5N~Ys5URvw8`pm(XE&YdE%(rMBc8~m zJ7LiE5UeL*@a#nxJbMe)hcI~dB@CWd5{BOZuzfG+x)csfOYBdue#8^3KVdk}Gk`GY z1`2kSU{@11!enlns_z=2Lw{aNSQt45!3eTH0No&=8%!8p@Iz~44%UU8$lT2 zofY+sBsyGgQtQ|#qC>oO^2c()J+&v_#Q(|CW*7rI*oTP^ZjC)fhwi5hMb>~_LXUsv)x z53HMa2;CoKnSke=LU)&7Tg5fHlOo;&qz2@2H;IQnTPWBf!r*z2V2cG?BG|oxeM0I( z4fhcp#_xWQVR~TN10J)};Q1h7p`F3BhXi|Au!iC~<0C|eTpkr{DPduhWe_|@7(5>* z47#%-&xb}^^TcfdJf9$*U{4Z;c%fjAk$BKOC3H^<_KaXX$#qTC@GQ|`d3ug8;=JA2zH6Me!f!ZE+xP7MlLT3-AjV466|HcRtvU9u(g7{ zBG@{?UL_36*=vNM&(;g}x?r`%_-*iXSf&HhHhRoX1Al|C(9U4mn}WS1SSxa!Ka3oM zV1Qt66HoN#JA~obo?R0}eeV(-a_KJmb6{bcaN;%+x{-ptN8+J|bi&{{lXyb+KGC7R zO@zVo1Hm>6_90=2w}mjo+bY<==z=lt&yR%eW5Ol`8BHCpd_r`{^HagL5r%QtF4zvj z5O1eop9xk?-0!=gi0uR3pSy(T=YkCv_Xmt74E?fOc-xC%cqL~`Ti+#5rhz>OzBMiA5Ck(n1ghBTsVbJ{~*w2Ff zBG|8j{U%s{ao2PWpd zg7A69+>hkOx1W%k{f%of)?cs-qc6bBnq1GT$+aKE`@p*g$JH(O@^tH~B~fQ zWuYtFBzf)6#h&PUcNwW_JbPPd+^UEeOY7O0GW)3=MCQGk&~@}`a+kyELie2*msMW> z+Eg$uHH7DWZ#>&NHPrNUog7xnV{W{g>)TJ6{o=B*%`+JK-erH78r*%a+9F0SZ&@2) zjpA0v)3pXL?{z(95&Ef%H!dza>oIrC>IqN$zi-s*>Mk}qcY9G^_^kKhySe=+mh~=P zKXtakxHX95Avbq?8yV)V7Y)7K-2T4AV_gxzyf^k@v~yTf!`yYPiRa_ybE(H%pUXVv zZX=q-3^lr}xnXYn%f-A`80PlJ_HMS2#ASR)Tw9CSCyP@nSWCfLc`Vwz#2&3ZW+%h5 z4PjxdYb#hg!B&j028s2-^W6)Rv-ZMsq|kK`x{iW%A}rmeXIf{%!iboAcm?Y!SU17C z3)VxhE5!Mqo;KX%hTAYjOmurRRAi)L`hJG0$*igc7|MWsq1N!Vx-t+4N z8|Ha-i0ig*!#!rFh4#g?5rT~rY!qRreY9X>2n(Z3%*c(jerTq^EeNfWxwLU*&! z-A#13w?kcHteWHeK0^)jM7(;1&ex$|ZV|fqg54_E0>TjQHo{QD?SkDQ*qwx7-Mou1 zHttXM3^8~oA2}8VW(L%*NpA))ef-NTJ-4XA3p<5z!_X)O~ zcn-B)Wt!TSzd&?hqzZzP73?3G{ets6cIpbDTPfI$o~QM~I{l*1y(HKw!H$-<5i)DA zpF;a$+RL8KJzuigV|ISzxkj+dq7!h5@t|Yplfm-Fx)d8IZeZ1`n`3cIH1BIg{DD>M zj68nm<>sEdc|~}RH0KqQd2SPIkI6fk?m(?X=Jx$Mk;|2$zN=~`dfHFu!zMzvN3arN zTy8Sw4I^v=nf9v4`)aT6?fl5CAUWR)md88)=<3!8&pr+9oP0S5$p1M*yggz({d1x2 z^7C3;JC@(JHhkRH8|IGhqV5)?&CM3`mhroNd&A5d;dt6c^SOt^H<-_E+}F)#x6TgE zyl=)VYdp{aOZu` zi&4jkpYAa?xA#5fZZ|f0%*}h1my0`oA9y;O3eKHw7VJa8wg|RWu#X4};cF z`9u3++77~^&4Xd={^HPzOk1Aaz9#l(XeXJrlX!*^34+fA+a=iNgpG+k)w#FbM2B4N zOSJ}xHQ;4c`fthiqS<`=00!?+3!5s~|j&3EgqQP6$?0#QTxx zP{U7zp@yFcgYFl?p!-#@>4gK+64zUJ{zi0QzY~Ude-H-UpM*j87h%x-O&D~g$C(nm z8o>G#NzVQeo|A;`q|lur3^kl440--5*crmc+p;l3jXOtIYCk(nAgoUv#r(pX;4Y`CaH(|&n4`ImVLc*ZSOBi(d2!k#^ zVbB#2Y(;8dTH-PY&w@mUaVSI>;uR(go<#_Qt|(!MSBx;|iVIePFyvB_FnC@`&IdtP zis-;f69&&RgrUB&gh4lxYjXYzI{!RYds}eauIgd`$n0Nt`KlJzfO_5WtYMgYeyO@wZ#x7q?=?N< zocVg>I{0eZB8&_q(W{s9~q@TpoX)V_V+JtO02Kh36XK`ITUk z$#*;Yg*|f?jD7Ee4fAT~i~`L208#^%>w$#fyXq>!CdR&36Bfqo8o{m=Y!G3{b1-3u zH$TH=Zzf zP9O}+%S6K9If*cMP8Mv6U~`5ArtvYlj_8o*RKk$w^@JhMT;jI|(}ZriU^4_OpYQzo zP{T~2n_NdE66|5Yc8p72-y=fzC}CJWmlB4) zevB~a9w!XCCj@&^uu!n41bbSrX9z-WIxd1iOA}U>aYp-xa#| z1WOm}eZe*f_JLrV1^ZC2xuW*DqOZ3I&#i=E-}EEGuzaS8Ts{`MEU9Mrl8oPNqz2?t zC%SQqk9)rs$vrYcB2&#NBu z;n59p_=+evETrV{Ly2;g0oB zUf#VNpPxNu5%M1F^~C^3_lwZ6-z>m$oEXp9aX+z|?6YN}u8MQ0*}(phxvaVO`Ty$W z=8pYu9_!}R^n1**#?##8=`X`vpFhoK_gd>8p0Bk?kA40Y>>t9ySa(veQ-q<9ZztC! zu?;&-ba=kOzl0%|i8*Xinf+})h4#g?GeY;DU}pt8N7#gna6xdB{Ub3K^mP`(qEi$` zv2W1%I-EP7l07+_NLXlRH!X$4LoQhfLk-yogYE*t!YEo0WEU)lU^xYQJ&!d{>`%nY zC3Lm&Cf98heU@8z*3WmoC;BCi(6tu23xzH(VThNHF!ak2QQrZQ3v`DB%P->f&Y!$r z3J6_6!3q(E<+CthsG*2pMFlG+SaHHI4kbK>ZJF(nk{+|uLi=J`DZxq$R)#R-QkJkV ziWUSH306+96QisFU;k2tuDoCc#rC>_&>bppelDoJqR=%H=lI(SR*85bUS+|m5Ek0h zO{*$cHNmP2HjtciN1io=?wqK-CegvOmSD97t3y~A^@@3}pJ2npb%T0BS6{G;1#2MK zy&{)}M29|WL>OwnM6hOsZ3pr7ud85%GFSciGgs!Q$&eT@u#t7Xk!P*JW_Jm;^IuM51I||l`FnD$r ztP5e$rLKTKI|iOzi4J*o6Rf*nJp}71Sk7p{;p=U);-(O9okk74gy&JB!}8gi=x}T? znCuTCULQ~A?!WZ)n9Vx0FQ#27SUOYf;{}^Q7;>3N7;>3J7<7{fgKi38&|OCubW;Vpo-lY$BMhF?3B&l!5Nsx4@SG*s z4T9ZI&a1=oMxsM~vk8N4j$pki1g0gf)6k70Y^a^xw6SEn8%D&OCKB=H3U;u(HOS1R zYwQuX`$7R0%vV1ju^F7`KJ z-JR*TiM)@uvNQ7B>Ui<(oQTm_jOTH3ZXCK9!jnDM9=al4|GL+UZWr}^RV#UY{Vq#n zUG2xX-68t-Wm7|xPCW0OB7O-`FS}-fdht8AEJ(<_gMA@9y9k!wWA1w0%R5i)F86nd z`uu0*yPkK)I$VdarJx3TuYG5ywnb)MH|9bSliiCC-=pOD_+U?aCxdCa{g`*_T-oVk6x z!31dM_&jOmb$y@keBC|>J?3sx)4dvPh%k=xl&7q0inV47t1|*xQ0#M%FQSz9V$+3O1Q|qK5Z~4lJE8*tO!BLR*Lq{j!xX^vfR4bE+*jZXXHV$AWz#*r$SZ5a(;R z3Eg(WQ2P$Sb`loGwU5tX|KQ{Pnb7SL>~q0(3syk`TGEA`J1q zCJgcR3D#IV_vss<+fNvBIUv}#gkf8LkTBG6h%m%EOc-=W1nVHqJsuUh?*#jvFy#3I zVaW5CU@gSG>h}nCTzH-ktWIiRTH^YL8h#YIp9K3^u-{0$(57J8FGAPNI|pUw5ABO- zzY5)Ng8eSoAB3S_{v<4nG6{n6E$tt?KmQ^+tPd@`^Agqt4FC5L$2u-MTG@U|j0c_n ze;(cAnZLbU+~c5sJm$vRQ9Ci7dkl0k)}`24;)Y!8bBWyJ+f!Z)_dLjHk41+_vERQQ zvy%~{S9SY|_eF822%Y~dT(^cx<8eXXx3Vy9XH1{D zW8XVEAkXY4w}1ch`qw?TayDj&;d%yh6QA>Z&hu%HWXyXOGp|2Rp`Bz}ieOm@8>;l5lW0svT=yC{_Q?OiwAzp65@(_kxPLpFMcwQ)Uc?HWSSbo6@5Qba|3RZ|P z)K^%rB7zkate9ZM1uG#~Q}Wyh)KF6BN(ok4urh*`73?Cx$`Lj$$Y`mAp}z8hU7y|R zGq)-BQ)pjIs~~h01*;@jWx{4=gbQvK*XpVeT^Q@C5{7=MMi_L}1*;)gO~GmjR$H*S zVjSuS-R#`96ZrVm6}ozY)feny!Pbjh8VFrO!5Rs631Jxb#)M(qn+VpFFl-|)B@CXI z3D!)o%LQvL*cF1c5UizOtpsZ=*j0IxFN1A_t}S60zjlP7uiFduv&f}`&~+5-0&&jw z#*xWu?<71s3)V%ju7XV$d3F=Jw}h^{(De|kr(nGV>n&Ix!TJi;g`9iGvUH`;?H9U! zLf2ogUxnuYp&KaJRf1hD*foM(E7%~x1`F1b>&ge+4%P&l?4sO&GqX=MWY~nV9?51e+^ZdvQ%EP3UeGY@T4Z z2sU4^TLoJn*lnV|+l20R!76#zsq7G797;ul#g7sC3Ef?UVPE(Tk;|Qe)f2hgC3NqS zTu|TLLYFS$Efl(m;#%7xp}U7L^x0y_ z_K;u?3wBu4Fi-T$Bf|4h!TOWmO z!JZRrnPATgwp_4)JbMlO@`BK<5NxGjFA^5r(qh&;bx-$8L^sN4Y*YUO_pqq%q+pfB zHRM$!7kt0JOjsCY83d~ZTO-(7!VvEj!Vs^ZxE5Siuz$pMVx5Ths$j1P)O=#cw&91RK?De zxD0}^|G@?Jl&9;1VCH?lh{t|k1yB21sARb`_3k@~Zm)=A9uV=0R!#1?zgBX_{$~a9 ztXU^f*D=n+zJ2TE(%xYQJ?7r4bI4;>hIs2M+fRIq*#E(Ru2p=$ihEBJ*aok!yV?+O zJM87s3c$P{@t8XfM?DtZrWI?y^OzfNeYEgo*5~%w_nyxE4bBf9bMLV^<}sT_80T`F zu=JRn5bQ_7;Q5naKMVGYV6Wz}NoDqn{S?|4(|#4Y-v~o}zYF#UVPO<42>ukTxj3)? zm(Z0G`}%(q9cuW8FywiXFz8MRcABu!aqZQ~{V4GKm*|ko8N#6Zk1*&4isQetL^n3h z^BiID4ASgOd|ApOSc+g-1^`4TrP)8Y?tr(wcE2BT3myBuMJ^?3CXaW4w=K7GLiaWKJu}#5FV9w17`HrW z_77g)sJhA7@8mlN@wSK>I(jukU$HifX&0v1SDwozZ@<9KiFhdulk3=To1y#1%f$Mv3)xu^;1>H_iUR`|NKKue7LPpI}Wr=6+A-6S=Vc z7mULd-m=uzmK3-AY4#6ZpTBS8?xPflb=j=kGp&s%+1-v5G@o6sliq%d zyH8WZ%xm*Q9CsT~+%VU-nECAPyQK6Dg2>uhm{v5+zGl|i&VGzri8T91CUcj8k}iVGHn$zubX3K^VzMnlKJfR zgB!PM%;4+xXAQ&LWuUtG?DlW9SeqrtXl@@?jPpgmw2J4kZ?)2bpqul#ruppVXFmJz-&x z(K-lL>v}s6@0X55hg>=l78)npI;Yv!L{I3t2-cOb&^X!FEzQ0rdWLqAY2DN8D`!0j z!+qz^Wj|jB_90=Ro!zvaY4(-J>qQu9?@d@}oNOC4%)WBZK55BWU%@5{&ntzlpJ4q3 ztCc4(jn^=M=+G|%1-pu{Fvw^dMZBwp?i#_a6>Jb;V}gt}n6POvdqFHqE)#600SI_-QG z+DWF3Nwcqsy3kHCZETu-=i+BqKTO`O;RgL3ChG3KkL;1{p1dYIt6F&d+haCv?k&?ghbC2)0tN z7X@orGBAzzStG$-BA%#k6=9)?Vp}ue`Er__gXgkZu(rZ;4bfqHwwACk$Y|3_*ms`G zD{0BuH6oYS$TAUHE7R7c*;gKKyzqQg=w1_Sy~xVWDxdt&PZYQ<|M4F?{ z?JM`ZPOxv%lI!*phFlI17FsLQ7EHCTJeP0NlC!&VCT9nU4!Il>>@Z;%_alVC^C)3q zkkP&)47%?H`+=}9$Y=#C*mvG%$I_Cs8bWtm==u{KYB)i3=(8UQgYGB7!XTslOc*?W zAuP<01-}x8c)t+_-HwXZAhAE8`(5b%5bRIE{vr(V{ub;X!onbzm(Wf!?cX%}n#iC#BiMh0g~rLYvuXA<(KED@Ogoom zUpWhIKA$0%EP|y7mX$CpyV*R}&4!QL1s>}PVBWKPthK{(c+4Wi`=GM@lvq0$yN3*n z{XcB5TdO2`cC;Vkmeb3nJAir5>4PBQKD?Dmzj3x#K1!SWF{(S|WCKVe}+ zf}nt41qCZ4SYg7j{uLn%xfB(wm|#UOw8ilLEG~3o@+H@m5W13rJxn}D+c2h;B08|r zgoTmH9KYwceT4p{T+%Oy*Hq{(B@DXD1ZzeZ@+>gn{4pA9Cz+;x`@UDyaJh)r zoG@&ct{^OoGBNjx5r*1Z5(Zr>!k}wS7<9eG`K2~O*H*9w;+#i2p_?sq?S<|sqC@Q+ zgs!7te-KafStp{y*s^ma&~+BNE`&kXl`z!TjWFo4rUs@Zes4n8UFdoUwnUuQxL2@B z70%Bkw3AHhDdP1KthZo&1nWx}`sGT(!YIoi*d%KBQm{V+TS|^Mu`Hb>`}1Jze-wgc z^NwX~`e1k_drDlB%Z61?{2zY(yuRq|X!?81_4!N0_kZWo*6|!5Jo|aaC~iFVyA#w~ z(#xft6K`M~4}H989^L#7{V?IK#k%Em9Y&2o%alA2tjU{Y)tQ$udY8WqA0dhSZ zx(PxzQLssZO(qP-t5XO=F73p1h}`5mYgF7X*AdT|G4p?OXeVPFrV?Ek>#i5I0n8+#5+{${5%nFhKN^$9OojJ-Px0S&J>=r1iL}74@E9t z3ASIb8%4aU$niM(Y_`zdL3CI?=MWt{ZxU=SVHk%r!EP3;J~>v0XK%ug=RDz=tKj*| z1UzpcI(W_}41IPhVenia*lmO%-tB_jA=sUQ-9;GkyqmB|wyR8AA=aaXL>ES?AXp^W zJ%pk5#e|`+mk@?{_X;+doKt}AKB7Y|_Y3v_VeouVu!jVjki&Lf;`SQ3JS=pN2==I8 zO9gw3FpR_FgrW8)2ty4|3KkLu&!+@?TCisXdseXL1p7jK_b($l)c(9+%Lzm6FAxUL z6@)>zQm_{Z!`Qw=7(7pseH*Y2ZrdAi448$9NEzFg0KN*p7^%O`4ILypyY>YLZ+xxMm`AYu_m9YeB9_ zB2RWr5{zAw1Y`dj9SqMf$ZFFy&ArC^mdUe2#^mVrP4n43u6`q~BPE{4o%bEX++*Ok zy?MJizV8|49_PLr&zmKl$BmipF*o1$J?0)~Z}OPCynWy?H*`xQW==EwzkTE}3&NOwEZB$SdhW{F*$aI=T5=C6Q17-_Jd%@1WTW72ZXPG$B7Poc0#Zp1^bCG^vlnJ z{UX?}g8e4gNpYUyccR1k{s&>`mp=(ZeN74_@9V#aZlaytw7&^MynhG_BUKQbBn&m2 zA`H6IghBT&VbJvy$Ny)9?mxm1?<`@6cTTWip6vv_Tvs2LoMj<8^mPhh$R#Ub&}AbG zx(fu$E?5t74nK#`YXEEVfoG|E05Qbbz3YIrIvB`5O zC3K|)D0gEIlL5puEsk zAPn2xiiCx+XLE6WsuIzq$GXad!E-9{L@rf`4xUvBgXcHI6Kt9|S5{4URu}An)WEdF zWghWr2whFW(yf(gwFnC%Vx9#j*j?rAOgxv06_T?$!gJ--$#w4w&$`01o?wT`aca6v z!L<59cd=maiDTGw!Vs?k@x=1kkTB>P5f(-j1i?&lKMdF&k_*^3f?Xowy(Z!{7P=;a zH6<)PuJ2O9&}aXVb8}!>s#t@}XGq#lVD|SAE$xK3UFOAT1z_HrdCWccb-7^8J!UNt z-#@qOvVB$Ur^Iy*I{&%0?)^4bhh@ zd(0x_vd%jv9Bp1>Z5ySlVrSxghI6ecv5$R2Jo_7uw$AjnBF2ekcE-%*$$1kvT05+ZVeUE7&gQf0+bPz<*Bz^_F-wX4 z-1VZnVeWbIZere^hPmfdd&Ki1W-BM(UNJ*0Zp^-hxwZ8%pWXAZy<=@k?B|~2?&r;W zWjrr(Xz%3E-(wbqF&{u!ddvnA2HjPJL3cG_&|O0qbk_D#OQ)1n4!VqtSU?T+^CD>@f(&Jpl5QbdF3N}u#@r0$v@g@j1kuc=)QDJM4 zxJ?7wBF>Ra5}uO>6e;87W{=saVFn83 z5eCm&2!pPv`5sTUt}GBG{vZg|Y4%k|+9lDbXR` zV}dY5)@_aBya@`Da?IhjKZrTRnIZ*7I z<|o$=!pJcQZV>15He*ZvV@%U;C0J!~{P(%g?G~(Rs_g{cXRQfKx3inJM|kcP>% z=K#O-n2i(K7t`)fwXZzSuY~Sv!qTmkY5NEZBVvy6yk)6(Jl(#1Lv--mFW3RX!r1qd zRQpF_eenEN=&F`KUx#haL7_V&*kQqr2zFGk?*yyT>iil)JIS=~iEc;+3I>T|(S71N z&JV;hJ@!0CSQxXHD%w91#|V9Wdwi^&63=Vjz}~E!`2V(!lNgxy-YQnfW8{fe7S86G z7|C?c3&#GBB64BRxdQ7~KY6_M-aX}ZPUxzO`t0w~>(ft;)Z1J>~4*SzE_xi$86nmYxwVxPY)QX4wgYE(c-5 zVqH$cpvy%V;^h`Bl^jFiT-v@tHgIBpB3>Thc_CrL;~Mf32G4wip@#f|6(9_r1qp*^ zA;Ag@Rz$F(gdvw=gdvyWgh5w=FzC9A{pEkf@jyx8S&A@>UunXSOBun+3U-lT#Iq0=+9b&h1SZn>cX?O(A5#Fu3+^9t1sBag6$qN{;$l;&d$!x?g@0S2f9M$&CBd2ZGQ;ZQsR$#SXx-wvRD)YgwbA>X|VeFf7u4h z=CC<1V9N<3&gF$+a|L1ORt(rm!mzorFmhi-7`j2i&4?RqWIk;B@;u(?ja)(zNt!pLF$fNda*I5!N~M#9Kp0e;O#U|{fKsQrxzhsVpjR^dW4A=&N&7A|?-v*|_>7qP$33R&(quxgeqs(>_hHiIZ z=tc)@kAUqNu>a26b?$owx}_R(z}X5p?A_4S=Nk5Du&gWdkF<>ou5~onrCijnJ zZWG44G$*dgKL3Ipo8}tpb6H1g?zS%b&-`r9s>WuAD#VIbC1PnbzmYuxazRFnd=DY>3S5!j^5YHF9Xxc1)1p zAr0NCh|=$416}hw6V|Kgj%(=ZZDdS?)#p^lR&0{W)}Wfr6BDan!<>-zTT{fUsV-vHBeAX$?E|yiaei98Ouq zbVk6=6gIBXoh1yLX9w&YVc0xZ7&gxf*k6QU^L%00yg(Sb3xzROyhzxdRk>U&Y*@ug zVZ<_}jJ=8Fl0bK9z%C2e<-+iHg)rj0QW(0cgrU1SVAlx4=C#7Gd2cXAxh~LMAFvyQ z5zCFjxHoZ=Fyj1cz-|uMfT@bKwfFznyd}`x8nD~NCi&eaI@s;Ph~;kqyF=KpoL1WY zE)0Kn3M*Tx^c_;d(A_Nz-9G|$k1%ZhQy4by4cL7Fn@`qSWApw%*E`Vl3)p*;=fd;0 zG*0}H`vU=cP#8HpBn+Dm3qvhca}Bv z_uoJ_z38Z;?*iTT0sBGN=p0_!Y~H(Uk@UVU(ZPNU*iQlbIbgp8>{nsKsyKfWMw$H{ zu;SQkpuHa?zh1%h*hGPD;(!eeY)%sBCJoqR!YG%?g;6e32qX6?12$E_rViLN!iaNP zVZ`!KRDi(2y8AXjI!=841Y@rL$|arbR(82(&qirRzZFP0-MVO?2y3bvZ5o0fdN~t z!5By7%$9Gk{I@dyNZSg+@VBBc&SNVDY-M3(%Ue-AHD~^hm-TkR{qR)+n}Y(je_(TP zpj%ZK<+579Ru@L@YXoe~fUPBrSk@MXzjXq(Zou}JXCSe;UZ7iF7;$bOjJC9)F#J6y zv7p;1&`m#2k+$}BMl2f#x+O$MESm(nL4j`5K(|@IhKfyc*gVh;3D_0^+cIF|Bo^vn zt3bDPz_t-ae%lJ89<~#PZhK+q{ur866ElgNGEkfe&+|e3j+4L#6o_HEZkYvw9g9zn->M_;(*;N zYaq$540KNgx=RAxpBL%O4}X^iy32&2yIdIMd4({_^U8o-B@BO83&Uo92S8W-%Q>>9 z6y2tQF0U!A+uS>`XH9KQrvZ2-oxMQ{kmpf2D-I^IL~b4USGF)LlEc7LAm_e zsF&4p*wuDZvRN;S8yi0A>(>9;VD;E;Zm?{ga@wttZ++e8Ee+jj1eSi^+F+~K*li7# z5&rgU_^Y#n8s%GOw>MbpFvuTB+wBd1gUA25_498Po22?%Ja!D7f2)`OKN_rlto7bD3nsjSd(E27eu>rB&)(CpQ)f$7cJN z)(=%$%k|NUOHm~>3 z0^R2U`%l28n4)WcUj({!gEe*z~(o?hE_Je z6^6h62JAaw#PYo`Z2l0i{{`&Ffc+G(p9A(wzOeQK=%}MfgyC;eVdy3kMq8R(7`iEh;cv=-O%<@IBNp5zo<_|-b4FprIa9!94%jRKn>Apw2_w$gg%RgP^Au_G<6(|KH)p`+5=Ja@3nP{v zf_o?P1iE=KN*0g0A__(p598wv`(;R|BNqt29_f*qmMF>|nD#q#=?2@ViX;D8Er^BhuwxFVJ}*TBA`G_4$~YwrhWH=$Z{XbsuXrSUv8w z8$MRA<+N^M^*N(;8g|yKeRjRT&IXCq=jhgN*s1&2u)(s69Vn z9a*HU{dk3KALxE=T;pde==uiN>VIt5+%pTSZHETSe=aMwW59+ABllqe+esMyzE~nV z_Hr2>=tcx=WWaU~*jiI`9m_7FBllf}k>4m`=ynr^ZufwV7KY6|gkf{ffbA8qy@hf7 z?IR3-+XUq~Sum$C`K(3SyiM#I_}edFPt4x6%~^u$`27Q$s|LCQL`NMRD2y^YC}0N* z!{#Bvuz6^}4hz`f!l?Hn0(PV@VmT^cM+fYffE^pK;{rA&U}FPze84W0dlck;LZCY_ zUDc<*g`wXk3frKf^8xbQcRFmNH?UEv`PYEW8jPWD4s`Fyn4SD? z33SVdKXkW>4tAR`^1D4?e+$?h0lP@z#NXco-JJpZUtset(UJSz!ieP`0lP;SHvcIM zoA(Oi7`rcE_Y1?{xPUz%41W&_!{$T6&^;WmM}%SX(SSWB41bRcqdeDXEedHrH==t& zbj0$cFl;^*u%`p|Ou(KE*zuABu{;;(o)<=bF9hsGVfcGV7&c!PhVEYhdqo&FUloSU z*M#j+jNh)2u@t)3MTfsPgrVCYxNdy2p{tL3-fFNMA#L#OfW0G(Sl$(e&G&?n-@k>S zdtVs-hU}W7_T%US(c$kyVc7gAU>^%3&QFA4^V5KRCX6^g7lyz82t)Tpz`hhloL>d( zYhn2NMi@3v4(2Yt6&?BgR~Y`j6Nc`4VZ`!7!2Tx;e?JPt=1&3pIbgpC!{4vM$nUoX ztKVb#y}^pK5oa+b@3kK*u>S;eNM8i(j^Ldg4+U%@iG|!J7Dk+t2tzk%z$Oz$c}^}2 zn^OoQmMH@^l`w2hEexBB%KA}s(*(L{12$d2_Ln(1{7oO|E*xE?&D+upf$k3hn=xQB z1#D(vSp0v+7#U`=LB@BOa3&Z9-!mv4SgKe5> znr*%Y%YQE`IDfzv2-t!FTPR=)2W*jmEgG;N<^2qlOE=NMx;I#ToYkYja&e*S8L(c$ zsK4IAsE0nn@YgqB{e)q2b6FRL&BX%U;=<4^Aq;=#HT(@vWyZFo=&;#eSXl{|3fO>P zt=YhUEiE>QWq>eZ`8t^QUMA2j8?b@Gh-Eop#Ii{c%kqKlDVb--=5Z_LLih8-Yoa60 z6~rHGMPc|`Nf^47g`rzT7`j0L8!QZ)OUgPz*s6hUHDScEx-e|65wHV;Sk@FB{?-yk z-?O$bbn66c-GHqpjQrLQ*apIgb3JVNb%WK7OaG+-Mp)R`Z; zO#`-hz>W&ouppM50yaEgkIDMVaXGEDjR{`M4xzrBQ^+dE+U1Z>}c?I(;__7_Gh2M8-$nHR-J zi{$^bml^EppuHRz*gPm;2MZ&IUl;9+1y*$HnjP7{oej#n)pkhW@ASYX-?KFlhSG0h zIU~?r7UXbjpxZKFR|PhA3v?d^x`*T*H2F;y*nFg+tFM_pbWAR+_O^u0ry4fbua)az z4PE^@_3#GEGHmYGuvup#g4`c&yvL!gdos{X(Jz0}^E+{|c2+N!BZ3@mT)bVE_sH+2 zfb}}GYu&8f+u5MJTWv=MvD_4t=hY3H*&K2H*fZ~Wxf~Vvdm$+6Pa6K}Yod=1Y);i^ zpY`02Y3S;0@Ysq?oKvYb@_VgYQPk}m7kqY}^w|wCkbaM;K08tEZ`f_V(;|Na_C~VP zRhhn7F^hGE#5T&(Y8%_|-)h1YJHEm4pYi`wRX672x8!4i%0~WmLc>nITu*GU{CRwQ zF6;bZcP1Y>|4H+&v<({ntM&7gWOHzB|B3C}-{Ej_;`R45oYaV+9^a`IBhR|-w8ZM` z_Wu;v`E$i4tKwR-7T;Nk)z`_NnfB}N^Y|deUytdG#Ou!lonG0VsIpoA?AaBgoa%ng zO|1UebJBjjZG4&h*7N&3?bq8}y$?ID5ko!xzcg6Ajh)|Mb$=H$SiQ_HY_NKLU({eZ z>T&r;+AbDGn<|BsEmcumA`F|C3TxdmPuk0bajkW^u+f#jD}?P`u`7k`RI#gs5zE!W z_N{c+2pd+hYlRW#b;3qgy6c5u^9EtVE8UI4sH2;Nk;7kw?OxfuSs1!oDmF>=+5C(0 zytT#B_wiKhwraoiim~*!+r=)i{!JJ;-XVDBP#Y# zwNG6=Snab$;gM>e_V_@xPq~as`%5NU?BAdE^H7@o`_leWHGlV}{iSRB|4jP>(mwX@ zN&Cyx_WzOga~Dticc=YE6j`gL|*m$#%V-jw{P_k0c8jFWQVy0G3pIgXB&a-m(nB#dKvGHDa!{;}lF z`TY}#b3_h3Z8J*yr0ovrkw4L1544xhdv!9(^Ct0!&G!d)>R_Kp4vc%O-=(a#mU7|z z_k+aA`EM1`k;59&CK!j<`L@hyrR`J6p=^<&_)Hk>+U7aQZ$=rTz+REEh8-%5+^r2# z51&h%w9k3OCd}RuOP~C(v};(Op1FY9+Zko{ABlxA*vn!Q-4~)G_kN3Y?e9y`(dIXh z+$poKL`Oei?FDwLv_W$CT5M9c2Z@e0AC>1fVsm6Jj`F$(pLOa`3;zajf(b2B^OSz!?PIL!WvDiD4 zVRrn{mIlgkL>)~f=S|q_Qr3(yzn3^^FU`5{96s&k8R@fNL!=(?_k;N3SyQ`Sp-=85 zIncM-n3j6@pZKFZ7nO4dx*tVH+xHV z$Ag`xIgZ9ko2MS;6@QfVU!)#rXPZkbl;^LKA932d&KYAC5(_qe6Pxtq(RIh~qQhn} zwzI5Z6A7bzo)esNCKeqwClSUuC(2<`(a}C9>fV_H>_|CosfRBHcIsfiNvD2g-9Q(a{&2B*z2H#!K`CQ;SX5G{U$J zSWd>Y)RDEN;kmlgHks6I*&;nA#xl{IWSu{omhH@z@&pFtSq`HS@Bj8W|R zo&5eFHmUa!k{`^@RoI+SY!c_nQb*`!YUo;rYE>>X2W*xG%M#kW)h*|aF)Qa!wx5s6 z{HL_dDmE$4wPdafc9@hI=MEcxmMw2ZF`M{<%`S{KF-O4W4A@-4I5&PQW2uoh}C# zCkZ2mH>KW*#m4NkiRd^wV!h5-IR5%e9nmKrBC)_mNuLZ`Pjrm+M+R(7sXuh~&crgO zm9}}M9!88u#VSF6cCO?=o9Hfa!fcKcW^pn%GsQAFsP;Ai`$7D{EC#}#Qr;lB{ zGZxMr3rj4pMTBAVHOUY5M!>d}`a^fC9MiCMrT*wYr~A_SrMg%_sh_d8G6cuw$0# zjHS$fN?R}SSGGt|^cIFqyAQ-Qtj&?Z`iM<(?;Egw!m#<4lqbyQW?*7T}loANF`)^Kv=gqPro;eNa$l zZ%LcSrhSJZ$L(TrY@+*8VnMffptCtSY}$ANwnSj_pdgkd1Ks@6f8uYM?w!XU%=$>K z;}(!HCTtR*bZWo+?Nk>zgfx*n@h;K z5%z-|W3;6e0)Hn3gkR;>6|#(zn7^lrn=o z8TjihWrohikuW>{V3!3p{}eF$?htGa2y(Dz$k5p{WU!xQObfH)4|ZUX!)F1D=Fe^k z^6Mw#NMgBN+RI+K%Sqd>a_s~=TF#rWX{9fKSzChHoFmMxJ76~V1ba5f&#qI^eJ?uh z+1R=g*nBdEre9iVY%bOIa}~OYr2b$Z1hLq&wCHvZY}$QMbbAK=M#`KAI=lY|TQP{` z)4-pNy~|vTX^!kqwjxR2y%fY^_dc21-gIA0F(vvs4`v}=2q zt?7n+8N_1u5zq||;`}XOc1-V{iz985NG!0mq^x1x4?lLJCbas6Uv-2O!-n#;uE$~;!H6C@ek(BkQTpVdzK>B2K8wa|{&w^RMwCnh4D29kXm_4rzv#~zR#u~5>mdlRY+XQTGSw8`@ zb;ht`rA@$Wy$x)OAoo7KyN+`m@kjk_8Q8Qj1F_gS6}GJSgW1|G*t~(x#+c}Sl;Z(* zUazjp#o8x2d&ZA4+e*r8m+U5OA9l|_+uIAw)_CrmKb|&wehg;!Q(?Bo1!n6GVVg=> zm#s*OVr$6{ot<-FcFuv>y-%2p`C!&wV7BHGW^JMh;Y1nOYtc<8~x9cO=Z-LI%3DB0d9h;4}9}lqY0=B&{V!5Dur%jm6 zS;6dHB)>HmNFJLYBa$*x^_lRI`i%r_~a4Bnac8vggQF4db^TjZGF92+| zAir^aa!d5SZbXp7GSc^u`^Z3N_brIi?o+}(3Cd*?sRwlS`~&O_DVH)AN7{A{au_H& zY}%L)c6gvWILN`SEwE{Q&(QHH=riuHe@HAaYu7M)_84YkKA6o%!fZWsS>^EED)~R{ z?Gv5NRl=;#g4s9{HbRaE*w16Sj%6yjFN3bT94oK^!piI>ZM#StL}%wz*xNx4U(3BL zbUy{Uv!#8avo(9LgXFk{+59}r#=9_E#{yeCXxFRCoYlDTMP2MFWzDlKqXM>@u(Gnb zyD)5iDfJF}Ps$oLIkfbAc!0|K^r$>cjG}zMJb29Dag1-|2n|9w9-HD>xExSqEXELYC{Ml>LpS2F3q_cPB zm)T(2R+sYwy1B(B^>Ci#4zqbF`jC?(Kh}8Knhe;gjlMO<4zqU!!gi7T$}CRX$r1~3 zo+7MlsnRhfj5cBOSuk6lg}*^)e!=g`w;>+Z`IDT>o27IgDUe(2VbaVt7|4+zZGe{$T~{dbtHMTNa6 z<1?5Yo8+*NtSKE;m5V(O#PNQHlo>Y94A@!1MppiwmKS$_tz zcN8!;^Si_X8`izkUzz`uwzDNau8ZwiFOKPRL|3*f7R9;3uxab8U{>!#E1Nc7z`1{_ zMt<~}SLPpSyRf0le=GBkv|S{O zK4iT==1)w>A@7y>N7^nH9mj*c>lkM3h4IT7a^8g5nmx{&4~njAMVaE2IMLaA;b51x zDy%(sm|gqg?-KDxoB3AGN3cspht11`p|kt9=q?u>>Hg)=MKAH1^Y(o7H03~gPkwyAYe9T zhgqBFTxD$uX79s=nLpagRZ{P;tA%mT+t!AnyGC@3%WX`{y;gf>l78Y^u{knVciL>8 zZ^U?1?6ge&PkTG#d~}`IjCvwFjKzZ0y#$#7^;8 zTOyX5#2>oo{p^1YbO*~ZhD}>rNbWa_P1r5M=&x;U47ytb-A3J0;WXL=Hf^m6%+^@b zw@xnYjJA85#6q063&W;8i%^6mJHU{ zEgi7C#U{DG738qC*rZ(k5!mc0*Wu{y5gqKG!iaNgi4$h$IM}svZlpZ#6@Re%gkdxK zj)40`M}FgkVbiWxh~)v%Ve>&@=pG8#!@`K=fS|0W6NbM>#3pe*Dhz)wOM8L!?_Q*> z{agh*AgJ5N#2;}!E)1KS%Gib+o(OdIEE_g|>(iM7HlGxmu&0EP!_xtKCScD7>^Wh? z^1LwevwPk!n{y9AdLKM z?hu*`FWtT^E-4XQjUx<#}zZ6D(UkM{l8<)dujSDuv7MsNKjWFf}>^*`oTb~7s)~b9f z{)qFx!ie)bVfeE-E0{eO0z1EZXT8I0UJBMz+9z>-AH;cqj3=;Z_v*0ugV==qPZ&1s z9WUsXoT*4#`|*Iz?t!8EG4L0iAAS-Y_5QOk{QV+~d8q+X4=_9T!{!kk{(cpI%x7&F z=za@y_ADv>eit3(QXF5Tt-X$56A8oS#KO=`B8>dr=-#!zNkxah$%J9k-c?QCI=Seu zIYq#x4A@k{@He$E{7n4l-2Az*(9*egpGX=`sw#CevS z$B1P{u}NE+NfEhG$o_Fie2eWw!4zDtr=78ZZZ&)a=i;#?%qS-EhGEh;)<=_ZU=?3oyt^}EE` zJ+NudUU476=2W>?@QIu|VD@fFn2lS>{q7+}+S=P0tcT>VbN+bRY)%|z?<;}Xcc{Va zokt_byDRK|6=mJ?`23;v97a?+d)65?R{G?!()ALXu-?M3`Kr_r?7V*2czc}K>?1Z| zeT8Au=Febue-mcs4*c~Kf3WxGDbnUSEGD|)`QvG`dm+@{;(;z2UoRm#Y%VFRY(|tSLOH~xhODwFZF`MXC2y`~b zf^Nk?7v0xBbKx9tdwCM)5@S0V{#Fuy)WgccINyG|PP@%k*O&j4wpALsI_uGxFR!yf zfo^bvZ8F|f>*uNsR=2rYgJnV4Ql+)8!dlCZ61RJiuyf@;0yftWo0R98!qBZHjI~Dg zJPGw+@0(%%)1HeYmbJwnv8)rYb%k-ics*e|r?SkOJu?EcXF|)CswmbMo797S-~2(9 z&gOz(_B~tp+d%yBeIXZ1pA7p~qh04PV6%yibpp1Yk9s#9Y(t5Yy4^?^`E4u=oA#VH zHa7`$R@V62G|=rI^t+n{x_PBM@wd6?U_*p)97Ss@wh$fVvZXNmZ56Psg%OKA(}c}! zM91+rx0DOa*1^NfAIz?O>0j&}$;7#>#EEWpIrpR6PIS!AFCsda&0oVdmN5g&zSA1k zBhcA9JIG;s$$@#{{ymDcwV$V9_TFCDrBb)B59HnL#4_o|o&I2dlsI8K2pgSErR}5O zxZN?(S$^2Gz7=NIKxHeEq8KXvV8euQ-0ma{e|AlRZg`-x>waw7T3gs{(yp0f86xYf zU^f0F&JhwPIy)D`>|Otq=g7e3_eZrqwD)nc{_}Oo57xV1r%mQj&XhLDe2Bfz z3TC>UvzxTpTCuW4(mTB8%>QZ654N+E3pVXK4)&;wBVo5J-sx{-{!`jc?wfzM#{#qW zR*c9WPn&(84a}Ygf=wlT56t@AvgIv3S1#p(ZVP#b8_e#}4zFzfE_32Adrp|+VONQT z^{rQR%f{Qwv#k7$3Us?l{bAGk0&MOkHo1mfwnui{9t+IYW@B^rz-BM0KWzTcztbN! zM+Y{e^;vrax@nf~+TWg{BbL2{k%PT!2WHPu4b35>&Ay)mX6xBucCLcixB%81>$keM z@u3uZH}b2qeFC;`!1fE+{sB86UC4OffwUbay0S%z;_!eSAq;=^{y}UW zDLU9u!nmHc_dvpI%#O{Y#U|_+VO$?gBmFMAV@1b!{v6T4?3oH|9v9fWZP6la6WUAO zV_a_cLdyJsw2cv)tgU!UY%=b#cf;&F-eh6tF~$XB#UISZ*PM&(yB&vD{zi?>|7p)1 z_JN#p%1U>mS{TqNCrncYQE7dQ5QcKS^xj@8p1;B8<79_oZ&Jd8+7O ze-g$VlAYgSn~Oh~ea9dEP7C~9FJn_|+IzRKd8>@w#uekYKZ`%+rS|BNKhxgMU^Z`q z&C>&$uLbS%j6i4O0%EcIaj?&&9_Y_jko;ga)`x8;{aHD_IEynS2iRG{nD1FLD3`NE zhs|?@Vbk6Vj_%w*H;crAO}pNNMfX0>6Mw|=7h$xe^Mx^=W$#OX+4T{uf6(__ApVH+ zLSfh(CG*fQyUzuiy=RfO_G1k8br9!85({N^u`q0w!pQFuVd!j50A}w}fZZx>34hNn z)|or(U&63y?}I>hspJ5=EMS)lV~k>L9@f0J$Qi+GO)2cKpl+{_Sg6}8g<;dmWv?7i z+H9_w@seHN!feeDY*deIxxMYeZ2dfJH_06~t+W@Iy&n_SN7@VQ5jh9I?E5!icE1>A z??r{#^Kme{_APTUrtK=JTXgn(9L(16z@q2Qt`>i^vulK5b3PeQz`hO2I$D!>t@tCB z>x7LSUu1b5;pC9;>`Wt(k~Z6 z`>_eTSLz7%pyUUO)^+_=a-ghl7KTk5qrmJv)-Zd&2J_eUj$~rFMPfm>spO9C)<9?L zsIh6!LXAo_lecL7(i*Kq+v|^OQCkOCwk#IKZ4xIsyB`3v_e8;L9mLL+zrND0VJAs> z!tDEqMpib@5}PnvM@H^DZIKPMml^DKDHm+oyf2i z-5;=V!k9yfo+El7(Am9KY}(oinB51Wy*w!X$nPOxLyPg7JvWH%;XoIyvw9@Z+4+`O z>^tRQW93>DHgm9MVZflAXZ!Jv&c>gxiDaxnxjY)=XWx-CE{Bk|$3(~6;gNyA$3=(D zCmO6}X1u%N$$&i-u%`p|Ou(KE*mD7UK432d?8Sh+6tI^YEEgB$`L716m**=1d$qyp z=YZE5tZwu5fV~m0Hv{%oz?#q0*JF7*(48?uE~N>_QC;^=pgXxWnd`bOj_BIvyMfJb zgIL}RbTj>-Yk%_xY@Wlr)}4Du*KGgRjjIW{=TDB#Khk#m*7>KiZ{;0>==!YGsiUp^ zyODeTB)YF#9b3CU=9BGQ%)IahGG2n&93bo}8Bf4$jTt$Nk?-&=b0TSbU-H{yJSyzI zA*@@E{Gs;TVfL;-%FNctv-al{DVLqASVqXZmCK5KAUU9Wd&T^r_Be-Ex(`Lys?x;m zyEeFo@_~#oVM8Pbm|csORV+*Q%KvGPlm7D~iIa9Vi}-{6HDLB`7}}D3hXu^8VPTd# z`8_y!XMV7cB|q*J+$T9OkNHWUdt2%VoA$k{uusJwaegZPV4n%YrmZo9*?0+N_h?~u zPZ(zR_+Ykfjdu2Vki%C&c`hn#mvZ?}VAI|$LM&g14x3*F>?>jPgQxc`($;>AVe@Oz z!M+iOOyWT{wOof|vs|o5TYC=ZY`q)D%Kl=LxmdezfxrJs4(O(k_1ByO z?7k?>&SUudPW*B1u=k0>?0us!dp|qv?0fM?EI$at=7&3*hFXdb{V7C+7WF2DmIzdm|Wt7*}O2!>K118 zH?+$AW0?n|@A*yQM7N}@Gp6nS9_VaLi_I70T*aKgfW>pb6Z+47Rs5}AC$zJ?M=TRH zSY0=9z$OXUqzy(tp8ZbNV0C|!2kgS&n(ODrwNu^Zcfobzqk+xE0-LV|x+#JjrVQ9r z0h?>aTykEw-5V@_lJ?mg+tkZt>cHP3!8PnO4PEv(I{!%9+rhQy(sC_|?(E?DcF~Em z#y6lQ;3w!q|V_KWX zggv!Nw&LYt-?g}N{&?EHmHw0KDw}IATdJb?gXBOg+lVCy^8VM z^nLU1_PT}joT`&?9N9Z_U^Zqiv#GSrCb4k54+z@X?19d{!>Md}D~dTp2b)tEu|#to zbBT_z;@raUH;*uMW2F9I_AViq%~|1Z-oT%&>Bgqri=={oulZ&YoGu-y(r7y63j2=qS%_!pPy5zMbVlEQiUwFuEV){D;5pfxrGzX4q_w zvvOp_(nD;*dIqeQF#Op%V{+&%x}iCsw0$Cd2+Zym!|Xj@`0LZ~S05AgZLs=vK)--3 z)?js;iwA6p2CMs9vcc*$k2x}z&V+uc&OSc4Yj)CMU9C6P-N+NBgvXkT_Qso9OI&wqWN;d%@-^ zVv{z%ywo4e*4A@8*z@QxTMGxXu_E)4-#23ESJi9&HK>t8oed7ystuNHmH9{7Rue{f zE+O*;)cfk9W8QSdvDtEaAJXc=(_ig-A7RluW!4aX=u(ftxqQREhRs6 zHYWhH{9v<69|E)dV0H}zn?0~;&v2l#cc*X;u=U=Y8#nD!q^-T5hxO~*H5(}PfX=RS zIB$L={Wa_n@drCp{K2ffQ14qwy>p&EQ~aUZT6CNn?b&0P-GheN_b~2~i#TnAB|n(G zr<-}A6+~CIEEdH!k^^kpfbA!9CHS-Nn}FF`cbJX!VYdF7I4_o1MppTqF1f?(J5ONE zdlT8*h{~qbZCSDHf->7)7&h&Fl`!ikVD|mR#PY|$pRFxKw}a?lI|>_KR>NPS(o|aX&V;Uw7wjh_C6MvjYDAeY+2cgB)u0- zVu1}8Hnh^&T0hu*k~_@i+3`0*{80}hg<;d4jmPG(gL1&_b&JmC#L?|6{+O4F-Uq*n z=&-qKz(xsUjIx_BZ0;_MG8^4s_3P<90=8$s_G+*VQ|Yqp-C%WpLmS6fo$V9o_6^v6 z0oy-d2L$ZEfE^UD=Dopseg_A-Ljrcn%(-MIwB5SRKLzZx2Fq5;{3C6LHvHwkmH9{7 z4im;Tv%QCo>py#jk?W(w#pdww?ur{%%m0~>U*2QfbBCOdVCIkU%j(@aZNk=&F)hc+ z5fW#aO{MKfVU)|d;;(F}isC5I;qPc+=#CM_wY|-|!A43fu)cDQ&iW8rlTY2+^Kry^ zti%aBP8c?$xzRDAqb>Cb=9n|JldmDSz`siWfqn|6N_o0CdC!0i45?9;)W zu~42TNGx1ao+u2PGxX}T3A1MoV24N!#Bx&L&yGKA+V}i0A7bypg6$-+z^wj=^JIw= z-8Dg5Iz@EYJXIL;4A)7VFkANo8_}n;j_~)Vz~9o+N1{7Tbg(}M>~vwo`TTO7al&jq z1a^2(E;a^3XLBXod$H$Q$?puwVRQ~XZ3E?+3-*wlA7JMPy2oU^i_X@n(7&7+#JQ4` z88(-fW4df*R}^OjHf?hr9-}Rz{-5`uuZWfy`JLkac zdW;-y6n{fg>F3SfI|;Mz-GOb|JF6$uU%$LtZ8tUiZ3alc|Jq=6-OT~JC1AHktbcay zao!f_ZV%Ys0(M7(F<#G6{Jp{Ixo@>pdz>5P-D zVEL1r$L49AAL_Z!(t7E!+x%w`%e?`cthHdkZQd8?{@lo+9_RgmZpvf3E|+nE?ty@< z)S4La93BjGGtAl^OWtF?^`V9?voilk+rz@>vmOzK%}0gpJKkjRn6R>9j|c3DfITUU z^&xhx4YTifU{1%@`jLap!IQ&3|3TckX$K zg*aaj#x=UVe|~f^ezSFlu=C{lmOj~@m4!{QQMNsy@5y^GyO&(%52Woy$${&?mxN){ z)||j>Y>G`=TLH5*Cop^eb=k_GC|;I0VgC|_O`DVBy76LZFR+ILX7A@thf_NwGpwn$ODCX9aR zbzxlpy&vTFMxgs%*1lrX=BQykd**=K%MA9VF!I}B%T68FV|G6PX8jUuX!lN=_Zm>*e7m8mvpbf` zk-6^4fW6bO$@x25d$+;r{+<= z^?WdU7e4H4Id8)DY|Mw`Y+*mk7_4>4n?AE=kNnB@a)FHrm_7f%IqpAFp4hbeE3m1h z|AZYV_xNDz1iG6R?>aww7ce#l_wCe?-xrb}*Lc%PJ;1h6AUdCXxJLy``Urjgj18dq|uxd)^*)KoHA!fz7p~?PBv5X`irb1Dp0eujuUF zJZx0p@0P&dQ*u6He(7du^RT~2nZX7v*V&d}R!1=F=V32NUk>|0a)()ey;lw|ZT9Rj z%#}0sN_nEQ=h0#Iy^qW*+L#t*Yd~OgOIw0%utW~HecVG1 z|CJojJ=~{L2b)*Q3}$Q0MrBiJvu9#pwhjn(GE!;mx>t*(xlo{9P z{~6S&gMA<9Y^@5qA#$w1?E5WXuSxyE?48Uon}dhFDfPZ{PC0GE0`_h1{E7DRM0e@f zu9>|@37vhf9;~+?0d3d`%AsUwhq{gQV&C`c4qgyVYapxX7fa_ zDdrMm9?0bR2j0ya@#^O)JXwmEU2 zyCTp{66huk*kl3QyfuM3p|d#xK@~pzA4f7o6j4EH$nezfC9pIIq~eDSe$ia|g5ca=@PMnf*^3)s~Gn>~nSj)2V>u)6|(a|OD`0^Qt!?#VzmtkGZB>u;XG z=DY!$FJK=9v0QUx*L6F8U~_?h?HuH{V4y4J=sK2(0ygz=UF#MK{4E@?MH(!77@dEl zt(&aCTlt3pS+HKqmABb)Bto;k;?@mENWD*G=LqE7m<= zJ%n*?d_(3LU^d2tS%1yk=z~&bFuOj2HT(IT8ue$-%949e$q${qt8>??%xvxgX7{T& z-##w+!MaPEfZ21&oQrMFk^Fi|ez4xcIH%gPbDUFcUJ9Fi#3pNh>>W1f`ic(LPZ&05 zkn)7t+5?Ww=SA1*#)_h_XFFks_evkAou~6&nSbQ1NB-f*t<9HnyxY7t>_-_B!R#H^ zFq`vW?EPCXHeF20y3A>%ZE<1PwD)SlY|Sjp)>&~s+@95Ae71zdQnsQjiX{WqUl{$T z-8+L#zDzdWUbpyLN^HWG7Df&OgrQqT7}qj(eFU@TK4JDO8|(s^6CjplB^GiRC=A_l z0b5=evDms;n7vOGW@}%GbA`a4&25vzilT$96tI<9@2GD}mYl0GNGe4b0X7!kW)F=8}Nf`gz#F za_%6%!IB?sa8+T9U#69^hS@zZn7yY4wrH=MQG4COtZ#+wAmb32J=X(!L&}A*lzl%Q z<+7TT3%YycTm_q1{81092R2`ovc~35GNy&uygu_i_Rdt8eFrRYt|4*4)(qHM0b4s@ z>jZ3FVWUzD%A2i4g*E4X^XFmqE|W5UAZ_c3zp_P&Vtrw|R?OaO%XN-D{|@^}+6(M0 zsdt#I4=ThIUFk+2L5ag5S!ac?#vz9do^M9Yzu6^z^3IevZ_B@ z(+#t89L%oAU{RmEk(3!a>(541aax)9D6v1}3O26XS#Vbh+q zhhblZ&W)WJUO*2&0yOUVJYl`vx2I$-AnaoW4Tv3Yx-vpEm)vpxiy z;}$5=*51z0{VM%6%+~i{bGcPIZNj#Za-lr86^3p*VZ^ee94qLy7ai=6!iE;(HyhK! z>^Q>a4q}u1b`*wgs4)E5^Y+*r7U-gDgq=i3EW?E{-!r8gw=nC=VRnBLenY{ylL zw6(WS*vLS)bHH{HMx6HEZ2aviI@l;-*tBB>HeXPlwk{TbyNN&8?g1MujJC8#gY_#W zY-+?py3UXO6kOy)eII-;o2e^?b~EoGxvU>#_YBac1vali1p2 z#_aa~YsQg%WvvS5MtgTdnVqI>e~E=TKpQ8+?41FW%j+^H2eWY`%;uM9X9q}}umb~j zP{0lrR<=r|D0Qa~2-uOr@E5(G;iy2jt@N$fwD}j9)gR36ckEWxdp{XR z!lG;Rqa_E*>=Sae_1o^3&Aa~?2z?wTA&DuQo zXl+d}=l9*^{7|+m7RA37&;OZF4|%W5Khk!N#4>z5D$W(gx#OJP`9tkCIX~Dr6=ri! zj1_;8tU_TCD{m@7${k;AQ$ z1MD_ojNN7r#?ZG1x`Cp@rhWG>_4l{H=8RHi8!d7~3@8UDQf9yGh$!4PBkx9k71{?4E$V8H@-28R+f}*l&T&`vTql z0sG*PTsjl#pEjj6OM93t;aGq+P!DzOmf zW5Tf6bE+KCgf^J>7z>YFtdlVowzb*R-yq2WX6p`Nvx$zeqP<(P%;BZ&amf$eRnnKk zY+X9czT=^6c`J%10)O-M${%VkPsXkG{vMd!gNE6AzF_to-~Y29mT3O$Ny(4gp9X7vuUxnG!-8OOsj5-0gRD-4^q1_YbWi4L333mctUN#1O&EjC{W zbhArefbK=nF;2GczJSd+C`Z)ZcCqN^>5f zVpRIaS}zIMKZ(I_BqX zT>|G+dv^+KHMxGk-|ONJT{Je`b=j_Scq6cB_dAH?%|K`CIMKZ&Ibk}V)}x2MaOlVeV1KX9u=#Ic z=-v<52g0~kJ5G)xn7!K-7M=SqT{$PyUe@R~7k@*m+-+`@v7*gS!0bK=ZDSJ@;<|_TPYg*I_Fs>=>x`V#1I4K7; zq0I7LnSZ2hB4M<9n`#*x(DbR5&u2%|5sXXjvd_UfnA*>s}A=Jdjd}f{PQGYWD!{5xp$ZwW_%_tgEA-eXNH^NBxnH%T7?vv+@EbN;~QuVNG30-}R0D2#KJ zogb)&g+#}B%=&AX%|{Z;!hubj6UU~_zi|GGzEf}!@rUlw-bLEl&jB#I=3*|$-g7c4 zn@Zb-a@`19TKvI!$o&rLVbPPa<@Pu^hg&RV{&?EDiHDWghIQRkGLiSYY;jpJo1d+LmeL*KhovTR)c# z*ua1-7qI06wnD)A%$q}+kY7E{Cy(x$y&32>X16`jfyY{zuz*Y@xuGV1L17n-9Tethmd$iBhMF(3$7}wMG-J`JXa;(5cG`?Fr z$5Q4WX)@Esu5*av@{hEw zC%Upl(t9$5QGerK3fyjU*Ge~}NB&PcD=W5v*koL{kc=Z4yFIdKr%jkW-w69!>YZ3N zlvp^gyxgkf|1<7Ix3G0~>tk9ycl7&h&_Be3Xsr_IGCeSxiOhFvDd zEo@S`?trbmT4&p3jP~|&T{9blp|k5aSoB>6LnL?lGqXuOY#}=8Z%bj=+(r5c<~#j!k0OG0@o>NNgS^{XEQ`M<>pq;tw`V7;)|tu;Id(W3m1W7QJ(9L}1h2 z;fBp0Ap|_s*Elvh^=8vstz>Pw#$^ zazPh8e>6JK*;=X*l|Os-5oX_;%ech1{?OSyT3Fe=(_fkYl(v1u=Fst|u=n=Dn)i}hhd`y< zw_%fUM#lREZ2tzU+dLp(2L^1tDYMrJ^-#BYP@p?FV21?k(10Blu)_m(M8J*=*oYwa zTjuY&T#7@xX1g`ks^m}7u8(TOk{Q?iM+@WH!qz&JS)8_GM91~R($cQEm;AJ_vSqO- zjuo4<(O*OdyR@;!EStmTablA(?I^K{Zj9&{pN$npc^)r}e&PgS_&ZS;=dq)gP6^eH zmAnTV)uU_Hd=4ysvdllyc2eMPs^0k%6OPTiSLPpSJ6Uv`|Ll8rhL4vOwtk|l*ymDz z=xjUzv-h#k=A-*~r$`Q5Kb$Izx?Nnxe6ZUW&B?Tv8OPCi(#~Mf9NC{F7Ru~2VH|(< z3>EA)iIevEXR(RD(}i)39@Xs`f$l*WyJ6FwA0w7C1DnQty8WY}z}dh~<3I!7dO+EYUL+7Y4egr0rtU zu3^dVBC$!F7YD2qhCf?dkH1Sqhrde$c9}5z*?aw9p9S;mmy1o#f8X~m($?N*k;4^% zE_??=ptJkF%puu4A8fE3e=xh2p)a@hki))}+^M6hBnQ~l0lOw(*9s#R`;HaZ%(9+{ zSgs3f+8h!#Z;|5&Hiw+6VAIGr1ZH!Uqf_h3+sQJQ02?lK3yaR-$1a&Ix3^t%Qw289 zU%YFbty#dPJ%a=LtxsT6$_(8qGM->AXt=6Y%V#bM^)|R z1Gz5)vwK;vangstY)%Juf{bZl1A2GXEzI)czOTI-2Dac~oi<^%b_-_jRe{-apRk{# zTwwMdESP<_747BkCw1n)^{u_19o9d%_WeZS#NQRmbnVa9qT+8%;LnaRbW;SeOd{uZ zbZ-R5*quR{*>_rF^GAt=>-Qt1pP`pW-AJ(4Bo z8>J0WE{_EM?A|;!*9*oIkBZIVResBN?~Jpo*kgfiBFP<_cHV^9cP+wfZiaGsJcz}< z!vdRUF4~zN>?~pA{)G60Jt+*EzsP+jbWa63dkzen_U=X4F16^2cF{vBr5YA~BugxR}<$nQCc1@^o!Y}#5Xn0;p{Y<(F|z@7+>qZcF=&YSD^ z=&V23PI8Qq--}`s_L4AczAkaXY#ksrUk+?q--_;Efo==AAAs(aK(}V^B5mz`D>>Ml z2R2_7o3!iKgrR#qUn}%H`)k_hRtwlzorx+U75T z&3PJgLG>KYZryP9diXW4`CGt#57?uvh2|6d<&5aFij&*(tFws$HnA}JrAdUbCVx5U zlez!?h>Ul+?`z*FQRZONHmUgAZ9FRMeGp~ECKDZFc6-+cEE*F{E;cD^n@^_9rVt(F zV)u(-pAOE+wAU?ch?F(CPbvObmk>SkGgY9odxO|~S?ZSc1vbA6vu9;#6W{di%$@wE zmN?N_nZfpxu{V9at$l^r+7Xz&a~d{})DicDZBBe>PC0GUNPcCD6vec{uo;atrV}0a z6zyI3uw^9&VwqlSvUckZX%jG;vl?CHXX84UwHMeDW3z$wV+Hn_vwNSgE2REl_6)$tD!&IsM=Y~QEa;+j>GO(BbT)T5qKf5%?)g9M?HYD~ z*d)#~=jmE^v)JVPXXOct&K+ya+v$(Fnd2o+m^}jyv-Nf4_pDs^)2?TgIzl&_Fm$sE zLpO&o<~*$4(QU9wk+$~wgUuP(w0F31ez;$3!u~G(F8R$R{)lDnfXx%Id4=I`K4HW; zi?j*Y1;PA;9S`X2dyio@zdJ6~Oy1@X;#?qLkIQj{zd7YzdD-$-6bp(?*g^p-2WJD` zmKGKrW7BP9ECrii=2R({MH)8yRq^Lviw11oLHQFC;;ies1-kAH#yBcl>(O9!n>_>8 zD`33?cHOpJ=g=q6U9wQuy1s$##zVW-^$T>11?=qwySBM_Lzh3vxTm-Dt!4f|+LjO< zZD~niTqFD@{SwbASU=C0{e9^tV73+&X7^LeRt81UU*be(@2ci_SSrxHGA0|J(B|`A znSZ2hY0;G}lHRK#jPd+1$sOG?qN5&`6~=geEwKrU#-9VlCb29hjIn-q>3d-IUI3U~ z+ryrgwgmfG&K<p~+iWjuE?NI#^BUy7vc!VU)_B6~ITUQ(8m#ZJ_oi_jU~89Qn+D~%isV39 z4-&>$-`3i~Y+edx_n={;B+e1V_|5tyn9bY3>|Ix|7scQ3%HLJp@_*XZ!OVbb8TU?`E`V$TQ^|q2_w#_rL1Amx`atvgTwYVNG$6I{_Opq z_`6x!1Z*Se=V3MnNGuyjER^Sl!q9CLu#E$@Nx(J@*k%FSb+ICCUe=ojx*@_C@7tPZ z*nsX```bcnQf6BQY%5{-J6P(S@!4W>9KqfSbngYs)(c~E`(SLdwd8>AbMeQR??YL) zO_^;YHqqH!ENypN(ZRM8hQI9t_D5mZ+(8&NcNB(hXuyUE!{$x_8!ilgBLX&382)w^ zMm^Yj-C;IA!8qC0y2EUZAk3aAgpH7PMh?434zPdERHUuF55eX{L0`VB*hIIT^vSR{ z<-9^HqXL_^h>o%TuR-p+iA{8y_UX)>bMfw?qb-dN*dD_0XWvJQ%{@g&x$Gqj-QEG) zM;JEu6-Jyk&f+>}cR2?T%YI@LovoRLMep|6Uu@#<0Ab{JV89L%MlAMyk+8ibcl;e3 z*xX2Tj5&cZ zi*?%M8vQG&BVrjNHeq9hVRL$^KbU>@Ha3qBY`z@WJbl%!^Rwp^iRG2Rrd_`i=Lr%g z>_lPsv-_#&P7)pLWMQ<4Q-q;ART!}hl-#lTC(&W^G+`WLe-?&6J06&~nM&Fq%)WDd zWHElbOUCmsyWbDHTl~SId#0yLew531GA^fF&Ioj82J9?h)X~|(MpWgpZJ+#~_Wq1G z&k-H$+<=`Y41f0CEo}Zpbd<~a!q8nHY z$nmz<1NX{pJWre#Nu02Yg<3O z3!4W8xz8>4$2g8I7k|WZMZm5M*j2)q|FmawVWXv86X(@plQ^#l*tG$>E?^%>nc?qx z(Glkj!pP6os$la*(ZOyC*k6U=Zw_fM^gTC=4u7`<>{em;8`7^xTl?`2vv+(B%^y#j zJ-b@ANKxD-{?OSPNSHnIL+*Cp2^Kvod%MJf&gT4Kw&s%j{w6kQOLquE_jh6F?i7aZ z_*HYkw$~ACoxz=q9PAna-Cg33Snd{v?jHfWCt&{!*mIH}vD_Qz?h{7tHfII1dv&mb z8vR-8@W@4+w)-3Y>TF!V9thZj0edK54>wqjk1~6t!Rm278nDL#_ISXa2-uSW`>av# zYvd%V?Wu-t9YFg1bc5CNdnRDd2JE?jJs+^g7S4f9I5z9Ce0OBm?1jMQivfFf%C2qh z6vVRjB3kJmOtS(-xeMB|9)#Nays&eiL<=g7OH+f zQ2nNFxO7bZL%R>Ici#~o%%{C8jJV$uMr{8UHmr*2ePLXqd>{;)9}44I_#T0sBlC_W(W@#{IJY2%|i|5Jv7_3M2QggwZCy7DgNSreaf7^^$+lCcaH#_Zk_Up0zoc5cs(ViwszU#3~(TFYY*JGPJ?bl6j$W6S&X*d|H)_1Gp(`_0&>(#vE}`GY_p{OdTcYN{d#OOrTuzrdA}apjA_3f z+aJ#@z1_Up0DnfB|k<^6hWbEN%xY_q5R9NX0S zSG6(cxNu^Nr@yhkNZPN*x=`A$$C~%+u`ZbQ>#;77_Up0EpZ4pq=KXrC^QHZItn;S* zW~{8!?VVVTmHj?xzaDF^v|o=k@7H7PnfB|k_DK8nSi7hFdaQZB9&5L>UypUsw4Y<0 zB>zfVj*a8EZ(>WPePUfK?e|anlykpI%egn-3gI|gBC%y^wid5GJ5l;u-UgHMCTDe~E+tItl=qT_#1Rs#~Nm-`B7zh3Tn zzh3Ukr~P`lFPHYQl`RiU`}K0q`}J~PHtpBTeVMeM%YE+rtJ)~{6sR~uUoY+FcFF#_X}^DMf1R|yRNBY>+G#(x zA@Kd!}yQ3j?wDB#kPs%W0tZWmG8Xk_9=@I)jnlCJnh%Z zey6lwFZ*F>zh3r3(|*0|cTD^Bvfm-?*USEoX}@0f+o%0{*>9Kjn`KY=?iR4!h0$+} z7RIr>hcL?S-Yq+0;aX!)(P8uPtvYqEy+lWUwzn|Gzh|!AX_Hv?5nY)-p0<63QP(f7 znLp9q&SCqBjy5xTNT;sMe@ff_q9eZpgzZ0GRvajdSPl|Kc|NddHqf5K?v>4-tZ;W9v|1*gQ-a?eoViI&+}S_gS=)9a!aexY#_PVn+zW=8?ka zr}mR}wr6E?c4^nUS8TmjA+(n#?ei${N8KJRjB+_f7;WiTVf4|*38T!$2rH``#s=(o zVYJ;7gmJ8_v_>}W<#M9vD6^A<5$DOmI37+Bwo{egOY3#ULYbW!=(gOjQ^#@i@n)Uu zpvvFn;*a+EC-Fy{cvOyC+Qez1qfPu-7-QL+rCn1;6Rq2s1F@VgHp?o%GXi#|Fpj^Y zw(ImqA99xH$o*_##Bz=>%KBVkdbe`y__rC;ey}q3`=?Bjj9kJLuOJKij z(rJ_2FA$rs3j=nn*rcDoNOZL8+vS+1o&CH~XPg|<7mH1Fr7(`o|45v$H~Mt#?-H>| zxm+5s%K~@>S{VMW3D~s(`%(Hj zY+ffiV!2)zx*LRXJX|SbH`rA{EH{cxV!25eHvby1n}w0XEyA#QYrt+3hRxfBk;C7F zp}RvEW1GJRZ20Cy+S|5z4=p!H4y3-%*Zm|ioHFm`E57AL(_XO;p z!gkN$rR`o}*t|~|y6F1he$kaJZ|NSeFwTwRCvV#82;BpsgFPs0RORm>VdVF4z#b8X zzej~(^Ra+EE)1Jb1nfy+oOAxyHz(R&X2kN8=r}jNE9WEF(}8Z@jXV9JvslRQ8LE(~2y$)RjzklxQ5=t>!va}Iw~bg;Jq_I87~mq0HRlV#GAkyeF9pNTuLR3-Ri*N^VCcRP4BfYau}ArJ zj2SP-PCI-jbW;jn)b}Yc7kn>t+^hUYj1`!=hr~SigRn__KMIEKC&9Qk^RBpu3sZjs z4{Ox6y3Szg|Em(8`o2f{?$5$c#?hg|PwqkdB6ReL-NbbOZSnrX-VDsN0)v?2V zb8H?e?%S}JdP2JE=UTsh4_!(SG#>BZkEswz$KV5RCG=i@6onz|l1n3_lw= ztg*wE6bxT#M4x~)adb-whM!9dhRtOJL$|D8=$3QX@(x=;Fv?p|Fv@G{u#r1eLMrz` zY_249u$3LQieUJ%s>4=u*tTL0!RG2hho5T*hHf8m4Fp?L=&-q#VCeP}b0oTEj&A&b zN=W5)z~9AG~Q`b4f(Av@6AljC_>AxFP*EQ@xvA$E@dZIk)WqrZ0Io$DO1EHh54IQ?T zVACzEkZuw43-{|T7Hym5s!FAeu!(Mqp>|SvtiT#M<+XKe4iopGvDs9N71(2@P;>*TD2kR&p<#iGaU1x`N5e%DM1*7kF6AWE9rW3awLhs}P1p*u*7ci8@Y&8bCg#Cl^#`wN?yfsh6WhRuP3Wx1+S*;=qE z1ygf5x_P}??~ZWT(qjCf8|mmqIqVv-{*5o{Y_!l3`xwD^CgVG&&bAdgu3^UtMtQ%9{Vuxg zgbt>zfv~x~qw6l#F>HFD?d~9Ka<61ZhpErG(Cy^t4iNh#;@es1xDMaNVe2K=xA~M( zsZF~ky0*n}dmERqo)OzEVTR~uy9>5Yp&c(+R=?yF}R+ z&*-0*3STlSCZzV;+0Sx$>|azLg!LA&!!8r$!PIpQ>~cqUg~P5CjJ4ws(LZ5o51m9ZkIohIJZyAVvs|ug*fqjtW~YaAtzfKU*9n&8s?fH{9B z_zek*bvGs~-x?RXn-Z3H*Mi-gFgr~Dyv1SbCVQKFc`9t)n&{#fZcA9aH@!V!`S!4| zc}K$Ra8|G<8`>|6+AnYFg54=}#CMm&?snKyN$ge>HkTLu40exW^InG?CFUM%9@)}L zUle=n%fMD;R%6*B-6zUp&ADGN*0=`*<67;kwr0HS3%Un|E-QRlDA=rmJtP=@J}elz zM;xZcJH9;X=vK-*YdJo2k2$)>1*84uh_wUui__1V4Yvv^`+}}_GQZfF>FJfi2vq?ikcH zI=UvE>(+J6C-!o^OtsTPddZ3JWx=M`$|}DkbDbT=W~2QrS1ynH1+NI3%ptEj>@|nI z?y$GTJzVN+!?x9aGH&gC9m-FQ=y}Oe&(>3MBhbsyjcHWpF1|^h`A136EPRS-rL-Q zDfbE3DsAg#Yj>)fO%n4wHg6HWOf^#>-6Q6uEGL!97b14%3HA5rls9;|8PM^4DQr^S zSAzY|PtoUL1BaE%i}&7N3!CgUzH!*M9m_T?$92`WLYLVIA$=zp_YU_EYbtC{v0s9H z@7TOApS;R-%{^Roe<{notW)6UK)w=Mh6^~Ib8`&-y# zjXOq+BiPHrCQN-U3j0Tt$A0$?G4Ey;R!IK}9c#|^;`3?P<6^9gt2J49Y=r$%t}~c= zE+H#){|R5veI~{&x{1OU=84J)X1weRte#-(ZR!h#PVKYsWwAt;yKOFnw0Oeo*i`!w z(j;+x1p6Twe{E|e`OhUBo72Vft=Mdk=;HahVZ!oZRg~E%VRjgsjUBe6!_xaGdALP+ zO&p#2+a%`e_ga~Oa=SCP)=TQs&Sb0{E3Qd0J0YZ{gfH~Zr3GWYR%>yVTdY);adiF1 z+DYXYW)-@h`_;{s6*kc=Cm8eR@`7bWc`FEp?#uiJWZ6%~(Ta|)sbK7FR&v-kVqRjN zUs>p2m*kt+vY)Wq#GW0tieqzC!KM{;_N&;NQlF~{9pk8zvu9sj=&-qlV61WKvjW(z zqCD7|!X|6ux5F&Za(%+C6>~Cvt|e?vGE*Tn6O6vJwqV@vsd3oRL(Fp7Pu9)mLI+d# zkhm7zQ>@dO6%$emVKd7~rLvA-*j(3PEd^sO-e9O1FUJQ{e_IAyx^J~EvqK@Z66Iw% zsZ?4E#yqI@yWA64FVVFv?AhD;3A1D9)cqHj+RwwT$``do^^y;NGaS+eNqKp5)}qRW z3A4l4+{j^V9JX(=7RRx-b##|6u+kUB5c|^3(Y1G2ql4o-`54N9J-=|tJA+2?|{VdlRd$$gvUNSp9q>Tl`W=Fx;GjtLRo%;J6=2rDt z4}Koi);uc5z#86GtW}wv9#UuF3+wcb;@T<8RpI+cg^u;FtHZiE?A*M9$}!M>>igZ8 z=erA=)Y&Eu+f*>F$2Jp;wPSOK9Vqr1)Y%q7$J()_V9dMki_hj^TL~Sp-!RA;t6ZP3 z9**uJu@6CaOuM>muHCho(XY2?R5x3AU^U|&L{H%-{d2j&)jC)&p`(BH7L0c9^3+hYb{rdf8gAEHAB68RW3Rf-%O12u8cB|DVSmOnv4?y$lsLVZ#I? zhHV@+++iaeHc~Lg-zdTGWwc=ITbCJXm0tc|DzGs@mtP**DIsm^*ct0E?>iv2P1d!V zyrkM6+c`FG%pEJw4fx$mF!r3=3!7Xg_uSCRE!$)b+Ck_TtLpExm^aj(8b3D>_pBMm zJBsqM+G?rnBp7{pXTi|zlCXTqEo|)2HPKNHF5Rt#$IlsXg}u6Jp^MM?I{@dN1n+y4CVG?jJ-IxIzGnx3^B&od-PpUts7^DLb^ok z<%sX_k#lSI8RFO3{pZ4;)}JnH()e$_hIWKmvbr-!tc zVCV)-vop%|0-GXqnVlZeRKd_qbJ)6~9kSeFc<#s1ZJQg{eol9EGX$f5)`)#G?64u` zRJlC%ax;ZZVwmNy*@7`1)aOZ4ih9{m^iSA7VjRI%>0ccKY=&U$yS5PZ0-GaZVD6dg zuz7;*S;RMAFzRz3!6p{Etw&eK$KL9qL3Oizg-vwoZx@)C)Sdx9_Y*c@`wNE60|Y~N zpu?UK*U8M02MHbQV8JNw7w2!y4iP$Rb`*V(Ir&hbgB>OqzT7IV2hknw=#CJKdH#gs zDC}D`6<|uP7LD*)*XZTJRX}Ti1`erzF&uSI4bd_ zrxjnMqZ4L_8PmrkEMG^C*Iefc9rbyh zV5}YLZwy%jF6(2RqdXU2bL5uQjCt^UQ6B69!I;nH=T4Q&Bfd=NU@x?;_64@v`gOCz z+ttnbZ&)|`bECT1+v2*4zI(Iyj#&0m^~4&^n7&ZN4!cOO%-jpW3FJ(XJOG}KXW|UHe;#N#cG$JTChR)F@bh}Xn0sz;*o}f=^CrQF?`DVH zlCZps6!m#)!t5|MZ*$n^V$MQ$N^7%Hu1{DGhc$NC!AI0>Q+*ytdA~X~=ZpJ9nH3Y# z?IQLpCzZ+_g7G~o>ieS6-6?d;FBLIo!CH#G%beU>=wL^Qxd3*T@C85b7K}MWt()lX z5jy6Odj&&xpTq8V*s~ogA(i_XbC3GIMOX(hw=(yrXMJHii27tNrT&hBJ#+){c^T|7 zac>6pfQX&BPJO>Pdnxrj7Wr+uP+N@$SJzo)heCQ#l*cu*`Y!M+SB3vA;^-zjpJ}cw z)+4Tu9uhWTON+G%UmkXJ{hab1NpwAm*zE1mgxN9b`c~1GV2?SvyOQhMeEVyLLwelN zO&nlnl>0p8?Jw$jp`9Mm6OPR%9d^9v^OV`l+>YiE>onH>sg?osx$%$Vo56ZOel zaN6h;d_yq&eA8iX35LzL9j4|G%6muX*fYHA zFm(?Co9{Wg_Z_DGw-5V}6~%Xdz(xx;Ewqm%^*s-;31jSx^4QGmP)HvLKN({m3dUaJ zBf+xVVx{u2V2pQlACtBJ9&sO<{q8lr&A5)?6Hy*>>ohUf;peA9N1y*pF!l_eJM0UG z9VOOj#?hBT$Mx7(g5l@af?@L;!LX_RcQmZ6Xb0--TVa#^%Xfm|=l2QgZ_Qbx9};GV z(f#PKpB&bpm02nG7i_Ba7B<(>sn6!oRW_)*yb)qwfKJ^9gH`6&ZF4tq{l~aHqFJ>L zHd)+1<9^mW@i_@|$Q z_lf8e__CO=30vG@O9(d2!U{<}p9O0mbXl&dR2m9KUuxvA#)4sUNryENjCFG9 zDSWiI)f_)xa(wA5)+5?|bzzh7x8lK8YB@gEqcwz%Yne3#V~$))FxI1Hg5k^B4m+}U zwV#>m|y=<{Dxy$Z}7@-#9roy9kD^tHZiEth->$ zXY1sPTDi_(n+P5I@=XO}9#r4y#yq&0(7`qrj5%Zr!O(3f7!3w zfi)RYH=B7vHOuT!NIgY)SxzdIUV?GYLw%nzthdl%vyWioi}Lz9tfjcd!)8B6H%a`> zH09kYt}S5wg-!N0>hE#!Wq{DZ1`3AF+0LF%eO`rbE77*>O}7^1!A6OB0yfCesn2E5 zz3tc>ENoIQuZ!5RIeJ(nq;j8UkD{I{fvL|FVMBy3u%QkcCKxenBN+FrhC6J8!$vyH zd)8r;(BbE3!SHiT!tzTEtNxI-O_&{{jmA1`su;KE+Bki8m>8Srwn^r}eCjiYLfS6z zC0|NvQDyst^(c+f;wogH3hM5q7(%Gp?PUKB&4p*y{V%&2|)Ro7w3h?Ialc z$Xy288Rhn4y*;IWHDeChS=h|%^pJKDEWZUCnED=cuB&zxI<8mL|36!pmlpotlh6_0 zZh~cnZg;_^6>Pj<>@g<@#@=nBV3fCqV3UjT)I85VYq_Czl5QjQ{bZE4rznp;xZ;Rv zo9HGvI`wQlHYW=mHun;YzBE!?`@)WEQ|%{AJ(ov$Q$%^Nse;j$rU^!Qdke<8GTmWs zI@frklKArNg@qi_3}F*rW+p72=Vv9%Y@z#ml;z5O2{wIT-E228XQ6B8upwQ_HftqT5cL8(!m)XzVEA&BVAz}>`Wb!x2GQ=YqlHaaQ!y{Wj&XFyI&6~#)oqK-3B&7V z$2m5S6FT-G--~&7Z!0FGhsC|aEGOarV2Scz3k1XFGr}h91fj#`iGra!NicLLJM0v} zuz9M(P7{psPIuTDg6&qcx|D0J+hKNr`&=q_?}>KRCEt|sO|*u}yoEnO`bYxr}6 z?Tm80z^-w0*E;M7$L66!>Mn1(Xm@;3&zEOrDx~WiU!E6hPL`{}|5$Q#HwZ@i-6+_E z!sfps26Q($x@*Pyj?KnmJtBsi9hfn&~kqvzS|s~`a4tZU*9ft z><70zp}IU+!(nx^24W3B*ICSUu=mCK4jU%gHZvDOy2FX#PKO;k)YRqrq`awOjmzxx zknR#T*(YBh`aJAzM|Y3InmB#%UPrg1s84*k&(Ym482#)4!La$D!wwVc`?%WbtTYwz zalN@v*n~YKSXS74STJ;t2!`%a!O$%+$c&f!3%bV~-DzUnVsiyy6ZW{UNqJ8QhEC19 zup`7=2Yb@7xt*A^(5d@iu<_zL4)&BNkN9356RJ9@>wHSFLm@q#=z7&2&VN3WFgr+j z&n7HymO}en!eZSn11z_w-P;x#dwV|7#q5QI+4=bLqQlhReWK%U) zV0eGYDP!?gcH*L#EQg(?TG!1sbaa2WsavO>bEdpTd8KJzZb-_D>-S|RzGXM4+m}~_ zZmOLgQYM~B=9>0ZN4HU^`?Nl>`I^uX`|E;jVx~fRLooJaZwkie18)h&HT2tpVe=iq z*eARz822mQ6KrB_2$lB*!