diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b4a473021d6..96540d329fc 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -23,7 +23,7 @@ files: ^(?!(?:flashinfer/moe_ep/kernel_src/blackwell_bf16_rank_major/src/|csrc/c # excluded from all hooks so formatting/lint churn doesn't block syncs or alter # the source-distributed kernel payload. # The shim/ layers next to them are ours and stay checked. -# The Cake concat-MLA, MoE finalize, MXFP8 MegaMoE, selective-state, MoE all-reduce, Mamba, +# The Cake TRT-LLM MLA Blackwell generated sources, concat-MLA, MoE finalize, MXFP8 MegaMoE, selective-state, MoE all-reduce, Mamba, # warp-decode, and KDA CUDA/host trees, including generated-dispatch sources, # are byte-faithful exports. Their SHA-256 digests are checked by JIT manifests # and must not be rewritten. @@ -33,7 +33,7 @@ files: ^(?!(?:flashinfer/moe_ep/kernel_src/blackwell_bf16_rank_major/src/|csrc/c # benefit and obscure its history. # GDN CP generated CUDA and its source registry are also receipt-bound exports. # Preserve their bytes; the handwritten GDN adapters and tests stay checked. -exclude: ^(?:flashinfer/msa_ops/cute_dsl/sparse_decode_nvfp4_sm100\.py|csrc/(?:cake_dsv4/|cake_sampling/generated/|cake_minimax_h3_mxfp8_pre_attention/|cake_minimax_h3_nvfp4_pre_attention/|cake_minimax_h3_qkv_quantize_pack/|kda/flashkda_generated_.*|blackwell_msa/|concat_mla/|cake_all_gather_matmul/|cake_moe_finalize_allreduce_fusion/|cake_nvfp4_svdquant_gemm/|cake_selective_state_update/(?:cuda|host)/|cake_mamba_ssd_combined/generated/|cake_trtllm_moe_allreduce_fusion/|cake_fused_norm_combine/|fused_moe/warp_decode/generated/|kda/cake_flashkda_blackwell_evolution_.*\.cu)|flashinfer/(?:experimental/(?:cake_mxfp8_megamoe_ep16|sm110_gqa_decode|nvfp4_attention|nvfp4_mla_decode|balanced_gqa_decode|minimax_h3_varlen_attention|kimi_k3_fused_router)/csrc/|moe_ep/kernel_src/(?:sm100/cutedsl_megamoe|sm107/next_cutedsl_megamoe|sm90/pull_style_cutedsl_megakernel|sm120/swapab_cutedsl_megakernel)/src/)|csrc/gdn/gdn_cp/sm_103a/|csrc/generated/dcp_alltoall/|flashinfer/jit/cake_gdn_cp_generated\.py$) +exclude: ^(?:flashinfer/msa_ops/cute_dsl/sparse_decode_nvfp4_sm100\.py|csrc/(?:cake_trtllm_mla_blackwell/generated/|cake_dsv4/|cake_sampling/generated/|cake_minimax_h3_mxfp8_pre_attention/|cake_minimax_h3_nvfp4_pre_attention/|cake_minimax_h3_qkv_quantize_pack/|kda/flashkda_generated_.*|blackwell_msa/|concat_mla/|cake_all_gather_matmul/|cake_moe_finalize_allreduce_fusion/|cake_nvfp4_svdquant_gemm/|cake_selective_state_update/(?:cuda|host)/|cake_mamba_ssd_combined/generated/|cake_trtllm_moe_allreduce_fusion/|cake_fused_norm_combine/|fused_moe/warp_decode/generated/|kda/cake_flashkda_blackwell_evolution_.*\.cu)|flashinfer/(?:experimental/(?:cake_mxfp8_megamoe_ep16|sm110_gqa_decode|nvfp4_attention|nvfp4_mla_decode|balanced_gqa_decode|minimax_h3_varlen_attention|kimi_k3_fused_router)/csrc/|moe_ep/kernel_src/(?:sm100/cutedsl_megamoe|sm107/next_cutedsl_megamoe|sm90/pull_style_cutedsl_megakernel|sm120/swapab_cutedsl_megakernel)/src/)|csrc/gdn/gdn_cp/sm_103a/|csrc/generated/dcp_alltoall/|flashinfer/jit/cake_gdn_cp_generated\.py$) repos: # Standard hooks diff --git a/csrc/cake_trtllm_mla_blackwell/generated/cake_source_catalog.json b/csrc/cake_trtllm_mla_blackwell/generated/cake_source_catalog.json new file mode 100644 index 00000000000..a9047263f8c --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/cake_source_catalog.json @@ -0,0 +1,480 @@ +{ + "schema_version": 4, + "target_order": [ + "sm_100a_148", + "sm_103a_152" + ], + "targets": { + "sm_100a_148": { + "arch": "sm_100a", + "multi_processor_count": 148 + }, + "sm_103a_152": { + "arch": "sm_103a", + "multi_processor_count": 152 + } + }, + "domain_order": [ + "mla_bf16_vquarter", + "mla_bf16_vhalf", + "mla_bf16_unsplit", + "mla_bf16_clc", + "mla_bf16_tail", + "mla_fp8_tail", + "mla_fp8_p32_qk_l2", + "mla_fp8_page64_pdl", + "mla_bf16_native_split8_pdl" + ], + "domains": { + "mla_bf16_vquarter": { + "host_sources": { + "sm_100a_148": { + "path": "host/sm_100a_148/cake_mla_bf16_vquarter.cpp", + "sha256": "cb7d0b90025ce917788c81b2e9378a94f65fc6bb1d0f95755c7f21cff00a3c0c" + }, + "sm_103a_152": { + "path": "host/sm_103a_152/cake_mla_bf16_vquarter.cpp", + "sha256": "cb7d0b90025ce917788c81b2e9378a94f65fc6bb1d0f95755c7f21cff00a3c0c" + } + }, + "device_sources": { + "sm_100a_148": [ + { + "path": "device/sm_100a_148/cake_mla_decode_value_split_paged_bf16_ee41d3ceea.cu", + "sha256": "ef2824fa3a3221a6c850fa71b29fbfee134afbb679bbed54be41523553670629", + "module_ident": "mla_decode_value_split_paged_bf16_ee41d3ceea", + "compile_flags": [ + "--use_fast_math" + ] + } + ], + "sm_103a_152": [ + { + "path": "device/sm_103a_152/cake_mla_decode_value_split_paged_bf16_ee41d3ceea.cu", + "sha256": "24650a23aa99b16d99a91a7f75c21b1bf307d355c43dd4981bd2c3a7b3e2adca", + "module_ident": "mla_decode_value_split_paged_bf16_ee41d3ceea", + "compile_flags": [ + "--use_fast_math" + ] + } + ] + } + }, + "mla_bf16_vhalf": { + "host_sources": { + "sm_100a_148": { + "path": "host/sm_100a_148/cake_mla_bf16_vhalf.cpp", + "sha256": "f528a64812a56db2fd519f800b0141482ff5805fc3821e9331a65ebbf19ede93" + }, + "sm_103a_152": { + "path": "host/sm_103a_152/cake_mla_bf16_vhalf.cpp", + "sha256": "f528a64812a56db2fd519f800b0141482ff5805fc3821e9331a65ebbf19ede93" + } + }, + "device_sources": { + "sm_100a_148": [ + { + "path": "device/sm_100a_148/cake_mla_decode_value_split_paged_bf16_fa7b59f3a7.cu", + "sha256": "97e416a87513b6f461099bfe494cee7c740efc22a11a5d4d53f1890d067b7b89", + "module_ident": "mla_decode_value_split_paged_bf16_fa7b59f3a7", + "compile_flags": [ + "--use_fast_math" + ] + } + ], + "sm_103a_152": [ + { + "path": "device/sm_103a_152/cake_mla_decode_value_split_paged_bf16_fa7b59f3a7.cu", + "sha256": "34663b398c802fc40c0e7681730cd5b8e46ecdcb92925c91c228e2b1d49a3263", + "module_ident": "mla_decode_value_split_paged_bf16_fa7b59f3a7", + "compile_flags": [ + "--use_fast_math" + ] + } + ] + } + }, + "mla_bf16_unsplit": { + "host_sources": { + "sm_100a_148": { + "path": "host/sm_100a_148/cake_mla_bf16_unsplit.cpp", + "sha256": "6aa4b7154454e81122872ead8cec35cc0b9c967ae4c23e1ada7f59f7462b65c1" + }, + "sm_103a_152": { + "path": "host/sm_103a_152/cake_mla_bf16_unsplit.cpp", + "sha256": "6aa4b7154454e81122872ead8cec35cc0b9c967ae4c23e1ada7f59f7462b65c1" + } + }, + "device_sources": { + "sm_100a_148": [ + { + "path": "device/sm_100a_148/cake_mla_decode_value_split_paged_bf16_5c36fe8228.cu", + "sha256": "8da7fb8aa58762ceb403e3aed3f81c74dbe8e246ca4a9283b5e713aa436225df", + "module_ident": "mla_decode_value_split_paged_bf16_5c36fe8228", + "compile_flags": [ + "--use_fast_math" + ] + } + ], + "sm_103a_152": [ + { + "path": "device/sm_103a_152/cake_mla_decode_value_split_paged_bf16_5c36fe8228.cu", + "sha256": "842d57f47f3cc2de2f09ded4468922c1ebfb96ca64a62ff9267b7ffd67c1cb02", + "module_ident": "mla_decode_value_split_paged_bf16_5c36fe8228", + "compile_flags": [ + "--use_fast_math" + ] + } + ] + } + }, + "mla_bf16_clc": { + "host_sources": { + "sm_100a_148": { + "path": "host/sm_100a_148/cake_mla_bf16_clc.cpp", + "sha256": "b44d87332e053415851ae557ec358c2f58daf5c11deb33245a92bdd8a78ecc46" + }, + "sm_103a_152": { + "path": "host/sm_103a_152/cake_mla_bf16_clc.cpp", + "sha256": "755d8f66b2e668029f021f31479d46384ddac2433e9b0c273051d2b1fb0e7bb3" + } + }, + "device_sources": { + "sm_100a_148": [ + { + "path": "device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_30cfe17e1f.cu", + "sha256": "a99b0d6beff6b2fc95b4d7a2d047753af62e4dda00bc3ef451427c8d22a24780", + "module_ident": "mla_decode_exact_live_bf16_clc_30cfe17e1f", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_3b5bd1e042.cu", + "sha256": "d687d8d03a68358e30145bb243f0d257ac9b3cd90f37009c2a101129511fcc75", + "module_ident": "mla_decode_exact_live_bf16_clc_3b5bd1e042", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_7aa2c0c00a.cu", + "sha256": "d4b4e462478f3f64d7481a7ca099da9735e8c7d70ef6a8b924b5c056a8cafb41", + "module_ident": "mla_decode_exact_live_bf16_clc_7aa2c0c00a", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_8587a85908.cu", + "sha256": "94b82ebbc180e40a757fb4e9239801815cf255af20ceda1c47f4287e50388574", + "module_ident": "mla_decode_exact_live_bf16_clc_8587a85908", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_8ace9da7a5.cu", + "sha256": "a6cc7ee9dfacd13419e5477e0d051f25f1b09c358cb71564e841df660fa8e06f", + "module_ident": "mla_decode_exact_live_bf16_clc_8ace9da7a5", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_bbfb74c484.cu", + "sha256": "ea2ea6812a7344550f12775d6dcd32f6cdc771bd6c7cdfadf90e89e3780c2849", + "module_ident": "mla_decode_exact_live_bf16_clc_bbfb74c484", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_c278f217ca.cu", + "sha256": "bf3a6e8b9ced749e56828a9725fa8ff95f3c98a33034900363914f563266ec1e", + "module_ident": "mla_decode_exact_live_bf16_clc_c278f217ca", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_c6609aef39.cu", + "sha256": "185abe504cf6c6210c7598bebdd5bef93bd7a72655477f9ebcd84233ae6b3397", + "module_ident": "mla_decode_exact_live_bf16_clc_c6609aef39", + "compile_flags": [ + "--use_fast_math" + ] + } + ], + "sm_103a_152": [ + { + "path": "device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_5f26657c8f.cu", + "sha256": "ba00e5fb0b470d27498b419f3ff2c01907612afa00a3757346fe4a237325cd36", + "module_ident": "mla_decode_exact_live_bf16_clc_5f26657c8f", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_5f6f81e2ad.cu", + "sha256": "a00f04af56e93d88334880d895589e068cc2aecc245c9d25b83310b2579fcd6d", + "module_ident": "mla_decode_exact_live_bf16_clc_5f6f81e2ad", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_823f60e1b7.cu", + "sha256": "e9e55a21d92e50b5d9f4b5cbd0afbd22653b82dd765158c2bbff8863a77da465", + "module_ident": "mla_decode_exact_live_bf16_clc_823f60e1b7", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_a6a70d7004.cu", + "sha256": "73889077417d7f4e3fc3ded7fd9c5053656e4c89322155bcbadb60e34b6a5828", + "module_ident": "mla_decode_exact_live_bf16_clc_a6a70d7004", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_d003c22422.cu", + "sha256": "dd73c1f425777cf9ce9e76a5b301bd02742addaafd6d4599d538d1be0eb9a19b", + "module_ident": "mla_decode_exact_live_bf16_clc_d003c22422", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_d3892bb3e6.cu", + "sha256": "6cb918fe119bc4e1d8f8e9d8653aec5394ef0dde5507b6fe4919e8ca4968d232", + "module_ident": "mla_decode_exact_live_bf16_clc_d3892bb3e6", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_d56d2460cb.cu", + "sha256": "f3765268d02e74bc2188a1f2451f8d1d98f49d5d9d873a28fea8131ab6029a09", + "module_ident": "mla_decode_exact_live_bf16_clc_d56d2460cb", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_fb6e96f4d8.cu", + "sha256": "3a32a12574ce8510cdacc03b7ce71441497dce617fa4fbdbbae1b3e523310a04", + "module_ident": "mla_decode_exact_live_bf16_clc_fb6e96f4d8", + "compile_flags": [ + "--use_fast_math" + ] + } + ] + } + }, + "mla_bf16_tail": { + "host_sources": { + "sm_100a_148": { + "path": "host/sm_100a_148/cake_mla_bf16_tail.cpp", + "sha256": "045ba76c818fd3d849dfe939917afd391f938da2d6c959d7ed6176f9e83e4e89" + }, + "sm_103a_152": { + "path": "host/sm_103a_152/cake_mla_bf16_tail.cpp", + "sha256": "045ba76c818fd3d849dfe939917afd391f938da2d6c959d7ed6176f9e83e4e89" + } + }, + "device_sources": { + "sm_100a_148": [ + { + "path": "device/sm_100a_148/cake_full_abi_tail_bf16_83451b7317.cu", + "sha256": "9ebed191e81649d19b0b4c507b78758ebda2441c68bac8e0fcc52eeec3036386", + "module_ident": "full_abi_tail_bf16_83451b7317", + "compile_flags": [ + "--use_fast_math" + ] + } + ], + "sm_103a_152": [ + { + "path": "device/sm_103a_152/cake_full_abi_tail_bf16_83451b7317.cu", + "sha256": "c98b67d0ccfc468c1942885d6d1bfb69f161fdea74af1b67557d5043eb4f5962", + "module_ident": "full_abi_tail_bf16_83451b7317", + "compile_flags": [ + "--use_fast_math" + ] + } + ] + } + }, + "mla_fp8_tail": { + "host_sources": { + "sm_100a_148": { + "path": "host/sm_100a_148/cake_mla_fp8_tail.cpp", + "sha256": "8211124df7b910a33662895a31b3ae26bc0dc97d3c2cb46f8390603c9707a576" + }, + "sm_103a_152": { + "path": "host/sm_103a_152/cake_mla_fp8_tail.cpp", + "sha256": "8211124df7b910a33662895a31b3ae26bc0dc97d3c2cb46f8390603c9707a576" + } + }, + "device_sources": { + "sm_100a_148": [ + { + "path": "device/sm_100a_148/cake_full_abi_tail_fp8_26b36022eb.cu", + "sha256": "3e86a610b4a69f55a9ff79e32680e20a7f213a6b19d96cd8f3c3d374fc483752", + "module_ident": "full_abi_tail_fp8_26b36022eb", + "compile_flags": [ + "--use_fast_math" + ] + } + ], + "sm_103a_152": [ + { + "path": "device/sm_103a_152/cake_full_abi_tail_fp8_26b36022eb.cu", + "sha256": "b1a59925f7c575ef3aba9496e217fc714f5da60c378ba22c15e5e4fc5fde3ef2", + "module_ident": "full_abi_tail_fp8_26b36022eb", + "compile_flags": [ + "--use_fast_math" + ] + } + ] + } + }, + "mla_fp8_p32_qk_l2": { + "host_sources": { + "sm_100a_148": { + "path": "host/sm_100a_148/cake_mla_fp8_p32_qk_l2.cpp", + "sha256": "1843a1371b9392f5a45d4f2f252126bab2418d8bfc855f039406b16a9f4c55b8" + }, + "sm_103a_152": { + "path": "host/sm_103a_152/cake_mla_fp8_p32_qk_l2.cpp", + "sha256": "57d9d49712844df776da710343ce6ae3a9f7c95ff1e2d5cd4d43fb6020c37de4" + } + }, + "device_sources": { + "sm_100a_148": [ + { + "path": "device/sm_100a_148/cake_fp8_p32_exact_producer_58f57b1b1c.cu", + "sha256": "24d0f049b8c0b87e8e18f950f990fddf29b271130e5d381e12ff5438896d8354", + "module_ident": "fp8_p32_exact_producer_58f57b1b1c", + "compile_flags": [ + "--use_fast_math" + ] + } + ], + "sm_103a_152": [ + { + "path": "device/sm_103a_152/cake_fp8_p32_exact_producer_105fd84009.cu", + "sha256": "2f4402efaf50df389e9306b0fe15e9f3ded64c1b62642c56615b267e8759c794", + "module_ident": "fp8_p32_exact_producer_105fd84009", + "compile_flags": [ + "--use_fast_math" + ] + } + ] + } + }, + "mla_fp8_page64_pdl": { + "host_sources": { + "sm_100a_148": { + "path": "host/sm_100a_148/cake_mla_fp8_page64_pdl.cpp", + "sha256": "3416d5129322f7c6a265d1ef63fa74730b7317825bfdce28cbc0c48737925587" + }, + "sm_103a_152": { + "path": "host/sm_103a_152/cake_mla_fp8_page64_pdl.cpp", + "sha256": "93497913cdf83a5eb0acf257f96398f9204d9f7dfc36178bd6d9274b8cc56c1b" + } + }, + "device_sources": { + "sm_100a_148": [ + { + "path": "device/sm_100a_148/cake_page64_keeps_gmem_separate_reduce_cea01a87e9.cu", + "sha256": "0255711ffe766bd5499cda419f43ed3b0eaa16fca59a627ae5819b5354c75158", + "module_ident": "page64_keeps_gmem_separate_reduce_cea01a87e9", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_100a_148/cake_trtllm_mla_decode_fp8_06282d4d84.cu", + "sha256": "67433bd8e01c5ca25d7af9310de3e26be7c1ef22cd31f12c938248dddf99dd00", + "module_ident": "trtllm_mla_decode_fp8_06282d4d84", + "compile_flags": [ + "--use_fast_math" + ] + } + ], + "sm_103a_152": [ + { + "path": "device/sm_103a_152/cake_page64_keeps_gmem_separate_reduce_cea01a87e9.cu", + "sha256": "0255711ffe766bd5499cda419f43ed3b0eaa16fca59a627ae5819b5354c75158", + "module_ident": "page64_keeps_gmem_separate_reduce_cea01a87e9", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_103a_152/cake_trtllm_mla_decode_fp8_6fe1e0de09.cu", + "sha256": "386cc5bcdb054f452631e0c40890d8916b83800ab9eb437d3094f338592b4ae1", + "module_ident": "trtllm_mla_decode_fp8_6fe1e0de09", + "compile_flags": [ + "--use_fast_math" + ] + } + ] + } + }, + "mla_bf16_native_split8_pdl": { + "host_sources": { + "sm_100a_148": { + "path": "host/sm_100a_148/cake_mla_bf16_native_split8_pdl.cpp", + "sha256": "cb5d6ec3aa37b7210858109e8b09e3234051111319f828f5148a031e654118db" + }, + "sm_103a_152": { + "path": "host/sm_103a_152/cake_mla_bf16_native_split8_pdl.cpp", + "sha256": "cb5d6ec3aa37b7210858109e8b09e3234051111319f828f5148a031e654118db" + } + }, + "device_sources": { + "sm_100a_148": [ + { + "path": "device/sm_100a_148/cake_mla_decode_v32_paged_bf16_c80238a3c3.cu", + "sha256": "eaebfcd9403afe77276b88b1dde6224b66c9928edd179fdaa047bec67b5cc1fe", + "module_ident": "mla_decode_v32_paged_bf16_c80238a3c3", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_100a_148/cake_mla_decode_v32_reduce_1db619cce9.cu", + "sha256": "daefedd5f11ff4dd42473190d515e02d93d73d09915d37a1386ec10cb7620fc1", + "module_ident": "mla_decode_v32_reduce_1db619cce9", + "compile_flags": [ + "--use_fast_math" + ] + } + ], + "sm_103a_152": [ + { + "path": "device/sm_103a_152/cake_mla_decode_v32_paged_bf16_c80238a3c3.cu", + "sha256": "eaee2652ed19052ff265a1c05cba04db99c295582943e55cd889ad59a2ee7fe6", + "module_ident": "mla_decode_v32_paged_bf16_c80238a3c3", + "compile_flags": [ + "--use_fast_math" + ] + }, + { + "path": "device/sm_103a_152/cake_mla_decode_v32_reduce_1db619cce9.cu", + "sha256": "7a8ef76e7f8f798b11098f79a229efd8c998ab75f1810494f08d3191eaec3f48", + "module_ident": "mla_decode_v32_reduce_1db619cce9", + "compile_flags": [ + "--use_fast_math" + ] + } + ] + } + } + } +} diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_fp8_p32_exact_producer_58f57b1b1c.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_fp8_p32_exact_producer_58f57b1b1c.cu new file mode 100644 index 00000000000..a80ba3a139d --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_fp8_p32_exact_producer_58f57b1b1c.cu @@ -0,0 +1,3195 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_OFFSET 0 +#define NUM_Q_PIPE_STAGES 2 +#define NUM_KV_PIPE_STAGES 9 +#define NUM_PAGE_PIPE_STAGES 6 +#define NUM_ORDER_PIPE_STAGES 1 +#define SMEM_SMEM_Q0_OFF 0 +#define SMEM_SMEM_Q0_STAGE_BYTES 2048 +#define SMEM_SMEM_Q0_STRIDE 10240 +#define SMEM_SMEM_Q1_OFF 2048 +#define SMEM_SMEM_Q1_STAGE_BYTES 2048 +#define SMEM_SMEM_Q1_STRIDE 10240 +#define SMEM_SMEM_Q2_OFF 4096 +#define SMEM_SMEM_Q2_STAGE_BYTES 2048 +#define SMEM_SMEM_Q2_STRIDE 10240 +#define SMEM_SMEM_Q3_OFF 6144 +#define SMEM_SMEM_Q3_STAGE_BYTES 2048 +#define SMEM_SMEM_Q3_STRIDE 10240 +#define SMEM_SMEM_Q4_OFF 8192 +#define SMEM_SMEM_Q4_STAGE_BYTES 2048 +#define SMEM_SMEM_Q4_STRIDE 10240 +#define SMEM_SMEM_KV_OFF 20480 +#define SMEM_SMEM_KV_STAGE_BYTES 16384 +#define SMEM_SMEM_KV_STRIDE 16384 +#define SMEM_SMEM_V_OFF 20480 +#define SMEM_SMEM_V_STAGE_BYTES 16384 +#define SMEM_SMEM_V_STRIDE 16384 +#define SMEM_SMEM_P0_OFF 167936 +#define SMEM_SMEM_P0_STAGE_BYTES 2048 +#define SMEM_SMEM_P0_STRIDE 2048 +#define SMEM_SMEM_P1_OFF 169984 +#define SMEM_SMEM_P1_STAGE_BYTES 2048 +#define SMEM_SMEM_P1_STRIDE 2048 +#define SMEM_PAGE_OFFSETS_OFF 172032 +#define SMEM_PAGE_OFFSETS_STAGE_BYTES 128 +#define SMEM_PAGE_OFFSETS_STRIDE 128 +#define SMEM_EXCHANGE0_OFF 172800 +#define SMEM_EXCHANGE0_STAGE_BYTES 256 +#define SMEM_EXCHANGE0_STRIDE 256 +#define SMEM_EXCHANGE1_OFF 173056 +#define SMEM_EXCHANGE1_STAGE_BYTES 256 +#define SMEM_EXCHANGE1_STRIDE 256 +#define SMEM_EXCHANGE0_U32_OFF 172800 +#define SMEM_EXCHANGE0_U32_STAGE_BYTES 256 +#define SMEM_EXCHANGE0_U32_STRIDE 256 +#define SMEM_EXCHANGE1_U32_OFF 173056 +#define SMEM_EXCHANGE1_U32_STAGE_BYTES 256 +#define SMEM_EXCHANGE1_U32_STRIDE 256 +#define SMEM_CORR_STATS_REDUCE_OFF 173312 +#define SMEM_CORR_STATS_REDUCE_STAGE_BYTES 256 +#define SMEM_CORR_STATS_REDUCE_STRIDE 256 +#define SMEM_REDUCTION_TICKET_OFF 173696 +#define SMEM_REDUCTION_TICKET_STAGE_BYTES 4 +#define SMEM_REDUCTION_TICKET_STRIDE 4 +#define SMEM_SMEM_O_OFF 173824 +#define SMEM_SMEM_O_STAGE_BYTES 4096 +#define SMEM_SMEM_O_STRIDE 4096 +#define SMEM_TOTAL 178560 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7])); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void elect_commit2(int mbar_addr0, int mbar_addr1) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%1];\n\t" + "}\n" + :: "r"(mbar_addr0), "r"(mbar_addr1) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit2(int mbar_addr0, int mbar_addr1) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%1];\n\t" + :: "r"(mbar_addr0), "r"(mbar_addr1) : "memory"); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x8(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x8.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x8_wait(float* dst, int addr) { + tmem_ld_x8(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_fp8_p32_exact_producer(const __grid_constant__ CUtensorMap Q_map, const __grid_constant__ CUtensorMap K_map, const __grid_constant__ CUtensorMap V_map, int* __restrict__ page_table, int* __restrict__ seq_lens, __nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_stats, __nv_bfloat16* __restrict__ O, unsigned int* __restrict__ completion, int q_len, int page_table_stride, int max_num_ctas_q, int max_num_ctas_kv, float bmm1_scale_log2, float bmm2_scale) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem + 177920; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 16) + #define kv_full_addr (mbar_base + 32) + #define kv_empty_addr (mbar_base + 104) + #define s_full0_addr (mbar_base + 176) + #define s_full1_addr (mbar_base + 184) + #define _source_p_empty0_addr (mbar_base + 192) + #define _source_p_empty1_addr (mbar_base + 200) + #define o_full0_addr (mbar_base + 208) + #define o_full1_addr (mbar_base + 216) + #define page_full_addr (mbar_base + 224) + #define page_empty_addr (mbar_base + 272) + #define s_empty0_addr (mbar_base + 320) + #define scale_ready0_addr (mbar_base + 328) + #define stats_empty0_addr (mbar_base + 336) + #define s_empty1_addr (mbar_base + 344) + #define scale_ready1_addr (mbar_base + 352) + #define stats_empty1_addr (mbar_base + 360) + #define order_p01_0_addr (mbar_base + 368) + #define order_p01_1_addr (mbar_base + 376) + #define _source_p_ready0_addr (mbar_base + 384) + #define _source_p_ready1_addr (mbar_base + 392) + #define o_empty0_addr (mbar_base + 400) + #define o_empty1_addr (mbar_base + 408) + #define _source_barrier_reservation_addr (mbar_base + 416) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + uint8_t* smem_q0 = reinterpret_cast(smem_raw + 0); + const int smem_q0_addr = smem + 0; + uint8_t* smem_q1 = reinterpret_cast(smem_raw + 2048); + const int smem_q1_addr = smem + 2048; + uint8_t* smem_q2 = reinterpret_cast(smem_raw + 4096); + const int smem_q2_addr = smem + 4096; + uint8_t* smem_q3 = reinterpret_cast(smem_raw + 6144); + const int smem_q3_addr = smem + 6144; + uint8_t* smem_q4 = reinterpret_cast(smem_raw + 8192); + const int smem_q4_addr = smem + 8192; + uint8_t* smem_kv = reinterpret_cast(smem_raw + 20480); + const int smem_kv_addr = smem + 20480; + uint8_t* smem_v = reinterpret_cast(smem_raw + 20480); + const int smem_v_addr = smem + 20480; + uint8_t* smem_p0 = reinterpret_cast(smem_raw + 167936); + const int smem_p0_addr = smem + 167936; + uint8_t* smem_p1 = reinterpret_cast(smem_raw + 169984); + const int smem_p1_addr = smem + 169984; + int* page_offsets = reinterpret_cast(smem_raw + 172032); + const int page_offsets_addr = smem + 172032; + float* exchange0 = reinterpret_cast(smem_raw + 172800); + const int exchange0_addr = smem + 172800; + float* exchange1 = reinterpret_cast(smem_raw + 173056); + const int exchange1_addr = smem + 173056; + unsigned int* exchange0_u32 = reinterpret_cast(smem_raw + 172800); + const int exchange0_u32_addr = smem + 172800; + unsigned int* exchange1_u32 = reinterpret_cast(smem_raw + 173056); + const int exchange1_u32_addr = smem + 173056; + float* corr_stats_reduce = reinterpret_cast(smem_raw + 173312); + const int corr_stats_reduce_addr = smem + 173312; + int* reduction_ticket = reinterpret_cast(smem_raw + 173696); + const int reduction_ticket_addr = smem + 173696; + uint8_t* smem_o = reinterpret_cast(smem_raw + 173824); + const int smem_o_addr = smem + 173824; + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&Q_map))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&K_map))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&V_map))) : "memory"); + + // Mbarrier init (25 pipeline groups, 0 ordered-sequence groups, 64 barriers) + // Mbarriers at smem_raw[177920..178432) + + if (warp == 0) { + // --- pipeline 'q_pipe' --- + // q_full: 2 barriers, init_count=1 + // q_empty: 2 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 9 barriers, init_count=1 + // kv_empty: 9 barriers, init_count=1 + // s_full0: 1 barriers, init_count=1 + // s_full1: 1 barriers, init_count=1 + // _source_p_empty0: 1 barriers, init_count=1 + // _source_p_empty1: 1 barriers, init_count=1 + // o_full0: 1 barriers, init_count=1 + // o_full1: 1 barriers, init_count=1 + // --- pipeline 'page_pipe' --- + // page_full: 6 barriers, init_count=32 + // page_empty: 6 barriers, init_count=32 + // s_empty0: 1 barriers, init_count=128 + // scale_ready0: 1 barriers, init_count=128 + // stats_empty0: 1 barriers, init_count=128 + // s_empty1: 1 barriers, init_count=128 + // scale_ready1: 1 barriers, init_count=128 + // stats_empty1: 1 barriers, init_count=128 + // order_p01_0: 1 barriers, init_count=128 + // order_p01_1: 1 barriers, init_count=128 + // _source_p_ready0: 1 barriers, init_count=128 + // _source_p_ready1: 1 barriers, init_count=128 + // o_empty0: 1 barriers, init_count=128 + // o_empty1: 1 barriers, init_count=128 + // _source_barrier_reservation: 12 barriers, init_count=128 + // Warp-cooperative initialization in physical record order. + uint32_t _mbarrier_init_count_0_0 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_0) : "r"(lane), "n"(28), "r"((uint32_t)(1))); + mbarrier_init(smem + 177920 + lane * 8, _mbarrier_init_count_0_0); + uint32_t _mbarrier_init_count_0_32 = 128; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(32))); + mbarrier_init(smem + 178176 + lane * 8, _mbarrier_init_count_0_32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 178432); + if (warp == 0) { + int _tmem_hold = smem + 178432; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: softmax0 ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 176;"); + { // softmax0_main + int cta_q_s0 = blockIdx.x / 2; + int cta_kv_s0 = blockIdx.x % 2; + int batch_s0 = blockIdx.z; + int seq_len_s0 = seq_lens[batch_s0] - (q_len - 1 - cta_q_s0); + int num_ctas_s0 = 2; + int num_steps_s0 = 2; + const int warp_in_group = warp % 4; + const int wg_tid = (unsigned int)(warp_in_group * 32) + lane; + const int col_group = wg_tid % 4; + const int col_group_base = col_group * 4; + float row_max[4]; + float row_sum[4]; + row_max[0] = -3.4028235e+38f; + row_max[1] = -3.4028235e+38f; + row_max[2] = -3.4028235e+38f; + row_max[3] = -3.4028235e+38f; + row_sum[0] = 0.0f; + row_sum[1] = 0.0f; + row_sum[2] = 0.0f; + row_sum[3] = 0.0f; + uint32_t _amf_u_0 = __float_as_uint(-3.4028235e+38f); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + if (wg_tid < 16) { + exchange0_u32[wg_tid] = _amf_enc_0; + } + asm volatile("barrier.sync 8, 128;" ::: "memory"); + int order_phase = 1; + int order_stage = 0; + unsigned int _phase_s_full0_0 = 0; + unsigned int _phase_stats_empty0_0 = 1; + #pragma unroll 1 + for (int tile = 0; tile < num_steps_s0 * 2; tile += 2) { + mbarrier_wait(s_full0_addr, _phase_s_full0_0); + _phase_s_full0_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float scores[16]; + float scores_lo[8]; + float scores_hi[8]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&scores_lo[0])), "=r"(*reinterpret_cast(&scores_lo[1])), "=r"(*reinterpret_cast(&scores_lo[2])), "=r"(*reinterpret_cast(&scores_lo[3])), "=r"(*reinterpret_cast(&scores_lo[4])), "=r"(*reinterpret_cast(&scores_lo[5])), "=r"(*reinterpret_cast(&scores_lo[6])), "=r"(*reinterpret_cast(&scores_lo[7])) + : "r"(taddr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&scores_hi[0])), "=r"(*reinterpret_cast(&scores_hi[1])), "=r"(*reinterpret_cast(&scores_hi[2])), "=r"(*reinterpret_cast(&scores_hi[3])), "=r"(*reinterpret_cast(&scores_hi[4])), "=r"(*reinterpret_cast(&scores_hi[5])), "=r"(*reinterpret_cast(&scores_hi[6])), "=r"(*reinterpret_cast(&scores_hi[7])) + : "r"(taddr + 1048576)); + #pragma unroll + for (int value_idx = 0; value_idx < 8; value_idx++) { + scores[value_idx] = scores_lo[value_idx]; + scores[value_idx + 8] = scores_hi[value_idx]; + } + int tile_token_base = (cta_kv_s0 * (num_steps_s0 * 2) + tile) * 128; + int local_token_base = (unsigned int)(warp_in_group * 32) + lane / 4; + if (seq_len_s0 < tile_token_base + 128) { + if (seq_len_s0 <= tile_token_base + local_token_base) { + scores[0] = -3.4028235e+38f; + scores[1] = -3.4028235e+38f; + scores[4] = -3.4028235e+38f; + scores[5] = -3.4028235e+38f; + } + if (seq_len_s0 <= tile_token_base + local_token_base + 8) { + scores[2] = -3.4028235e+38f; + scores[3] = -3.4028235e+38f; + scores[6] = -3.4028235e+38f; + scores[7] = -3.4028235e+38f; + } + if (seq_len_s0 <= tile_token_base + local_token_base + 16) { + scores[8] = -3.4028235e+38f; + scores[9] = -3.4028235e+38f; + scores[12] = -3.4028235e+38f; + scores[13] = -3.4028235e+38f; + } + if (seq_len_s0 <= tile_token_base + local_token_base + 24) { + scores[10] = -3.4028235e+38f; + scores[11] = -3.4028235e+38f; + scores[14] = -3.4028235e+38f; + scores[15] = -3.4028235e+38f; + } + } + float tile_max[4]; + float _max_0 = max_noftz(scores[0], scores[2]); + float _max_1 = max_noftz(scores[8], scores[10]); + float _max_2 = max_noftz(_max_0, _max_1); + tile_max[0] = _max_2; + float _max_3 = max_noftz(scores[1], scores[3]); + float _max_4 = max_noftz(scores[9], scores[11]); + float _max_5 = max_noftz(_max_3, _max_4); + tile_max[1] = _max_5; + float _max_6 = max_noftz(scores[4], scores[6]); + float _max_7 = max_noftz(scores[12], scores[14]); + float _max_8 = max_noftz(_max_6, _max_7); + tile_max[2] = _max_8; + float _max_9 = max_noftz(scores[5], scores[7]); + float _max_10 = max_noftz(scores[13], scores[15]); + float _max_11 = max_noftz(_max_9, _max_10); + tile_max[3] = _max_11; + const int local_row_idx = lane / 4 % 4; + float pair01_left = ((local_row_idx % 2 == 0) ? tile_max[1] : tile_max[0]); + float pair01_right = ((local_row_idx % 2 == 0) ? tile_max[0] : tile_max[1]); + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, pair01_left, 4); + float _max_12 = max_noftz(_shfl_xor_0, pair01_right); + float pair01 = _max_12; + float pair23_left = ((local_row_idx % 2 == 0) ? tile_max[3] : tile_max[2]); + float pair23_right = ((local_row_idx % 2 == 0) ? tile_max[2] : tile_max[3]); + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, pair23_left, 4); + float _max_13 = max_noftz(_shfl_xor_1, pair23_right); + float pair23 = _max_13; + float group_left = ((local_row_idx < 2) ? pair23 : pair01); + float group_right = ((local_row_idx < 2) ? pair01 : pair23); + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, group_left, 8); + float _max_14 = max_noftz(_shfl_xor_2, group_right); + float reduced_max = _max_14; + uint32_t _amf_u_1 = __float_as_uint(reduced_max); + uint32_t _amf_mask_1 = -int32_t(_amf_u_1 >> 31) | 0x80000000u; + unsigned int _amf_enc_1 = _amf_u_1 ^ _amf_mask_1; + atomicMax(&exchange0_u32[col_group_base + local_row_idx], _amf_enc_1); + asm volatile("barrier.sync 8, 128;" ::: "memory"); + float old_max[4]; + float new_max[4]; + #pragma unroll + for (int stat = 0; stat < 4; stat++) { + old_max[stat] = row_max[stat]; + uint32_t _amf_u_2 = exchange0_u32[col_group_base + stat]; + uint32_t _amf_mask_2 = ((_amf_u_2 >> 31) - 1u) | 0x80000000u; + float _amf_dec_0 = __uint_as_float(_amf_u_2 ^ _amf_mask_2); + new_max[stat] = _amf_dec_0; + row_max[stat] = new_max[stat]; + } + mbarrier_wait(stats_empty0_addr, _phase_stats_empty0_0); + _phase_stats_empty0_0 ^= 1; + float local_stats[8]; + #pragma unroll + for (int stat_1 = 0; stat_1 < 4; stat_1++) { + local_stats[stat_1] = old_max[stat_1]; + local_stats[stat_1 + 4] = new_max[stat_1]; + } + tmem_st_x8_f32(taddr + 32, local_stats); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(scale_ready0_addr); + float safe_max0 = ((new_max[0] == -3.4028235e+38f) ? 0.0f : new_max[0]); + float safe_max1 = ((new_max[1] == -3.4028235e+38f) ? 0.0f : new_max[1]); + float safe_max2 = ((new_max[2] == -3.4028235e+38f) ? 0.0f : new_max[2]); + float safe_max3 = ((new_max[3] == -3.4028235e+38f) ? 0.0f : new_max[3]); + float2 _f2_0 = make_float2(bmm1_scale_log2, bmm1_scale_log2); + float2 _f2_1 = make_float2(-bmm1_scale_log2, -bmm1_scale_log2); + float2 _f2_2 = make_float2(8.8073549f, 8.8073549f); + float2 _f2_3 = make_float2(safe_max0, safe_max1); + float2 _f2_4 = make_float2(safe_max2, safe_max3); + float2 neg_max01 = fma_f32x2_rn_ftz(_f2_3, _f2_1, _f2_2); + float2 neg_max23 = fma_f32x2_rn_ftz(_f2_4, _f2_1, _f2_2); + float probabilities[16]; + float _fma_0 = __fmaf_rn(bmm1_scale_log2, scores[0], neg_max01.x); + probabilities[0] = _fma_0; + float _fma_1 = __fmaf_rn(bmm1_scale_log2, scores[1], neg_max01.y); + probabilities[1] = _fma_1; + float2 _f2_5 = make_float2(scores[2], scores[3]); + float2 _mul_f32x2_0; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_0) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_5)); + float2 scaled23 = _mul_f32x2_0; + float2 prob23 = add_f32x2(scaled23, neg_max01); + probabilities[2] = prob23.x; + probabilities[3] = prob23.y; + mbarrier_wait(order_p01_0_addr, order_phase); + float _exp2_0 = approx_exp2(probabilities[0]); + probabilities[0] = _exp2_0; + float2 _f2_6 = make_float2(scores[4], scores[5]); + float2 _mul_f32x2_1; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_1) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_6)); + float2 scaled45 = _mul_f32x2_1; + float2 prob45 = add_f32x2(scaled45, neg_max23); + probabilities[4] = prob45.x; + probabilities[5] = prob45.y; + float _exp2_1 = approx_exp2(probabilities[1]); + probabilities[1] = _exp2_1; + float _exp2_2 = approx_exp2(probabilities[2]); + probabilities[2] = _exp2_2; + float2 _f2_7 = make_float2(scores[6], scores[7]); + float2 _mul_f32x2_2; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_2) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_7)); + float2 scaled67 = _mul_f32x2_2; + float2 prob67 = add_f32x2(scaled67, neg_max23); + probabilities[6] = prob67.x; + probabilities[7] = prob67.y; + float _exp2_3 = approx_exp2(probabilities[3]); + probabilities[3] = _exp2_3; + float _exp2_4 = approx_exp2(probabilities[4]); + probabilities[4] = _exp2_4; + float2 _f2_8 = make_float2(scores[8], scores[9]); + float2 _mul_f32x2_3; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_3) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_8)); + float2 scaled89 = _mul_f32x2_3; + float2 prob89 = add_f32x2(scaled89, neg_max01); + probabilities[8] = prob89.x; + probabilities[9] = prob89.y; + float _exp2_5 = approx_exp2(probabilities[5]); + probabilities[5] = _exp2_5; + float _exp2_6 = approx_exp2(probabilities[6]); + probabilities[6] = _exp2_6; + float2 _f2_9 = make_float2(scores[10], scores[11]); + float2 _mul_f32x2_4; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_4) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_9)); + float2 scaled1011 = _mul_f32x2_4; + float2 prob1011 = add_f32x2(scaled1011, neg_max01); + probabilities[10] = prob1011.x; + probabilities[11] = prob1011.y; + float _exp2_7 = approx_exp2(probabilities[7]); + probabilities[7] = _exp2_7; + float _exp2_8 = approx_exp2(probabilities[8]); + probabilities[8] = _exp2_8; + float2 _f2_10 = make_float2(scores[12], scores[13]); + float2 _mul_f32x2_5; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_5) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_10)); + float2 scaled1213 = _mul_f32x2_5; + float2 prob1213 = add_f32x2(scaled1213, neg_max23); + probabilities[12] = prob1213.x; + probabilities[13] = prob1213.y; + float _exp2_9 = approx_exp2(probabilities[9]); + probabilities[9] = _exp2_9; + mbarrier_arrive(order_p01_1_addr); + order_phase ^= 1; + float _exp2_10 = approx_exp2(probabilities[10]); + probabilities[10] = _exp2_10; + float2 _f2_11 = make_float2(scores[14], scores[15]); + float2 _mul_f32x2_6; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_6) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_11)); + float2 scaled1415 = _mul_f32x2_6; + float2 prob1415 = add_f32x2(scaled1415, neg_max23); + probabilities[14] = prob1415.x; + probabilities[15] = prob1415.y; + float _exp2_11 = approx_exp2(probabilities[11]); + probabilities[11] = _exp2_11; + unsigned int regs_p[4]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities[0]), "f"(probabilities[1]), + "f"(probabilities[2]), "f"(probabilities[3])); + regs_p[0] = _packed; + } + #pragma unroll + for (int value_idx_1 = 12; value_idx_1 < 16; value_idx_1++) { + float _exp2_12 = approx_exp2(probabilities[value_idx_1]); + probabilities[value_idx_1] = _exp2_12; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities[4]), "f"(probabilities[5]), + "f"(probabilities[6]), "f"(probabilities[7])); + regs_p[1] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities[8]), "f"(probabilities[9]), + "f"(probabilities[10]), "f"(probabilities[11])); + regs_p[2] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities[12]), "f"(probabilities[13]), + "f"(probabilities[14]), "f"(probabilities[15])); + regs_p[3] = _packed; + } + const int matrix_idx = lane / 8; + const int matrix_row_idx = matrix_idx % 2; + const int matrix_col_idx = matrix_idx / 2; + const int thread_row_idx = lane % 8; + const int segment_col_idx = warp_in_group * 2 + matrix_col_idx ^ thread_row_idx; + const int stsm_offset = (matrix_row_idx * 8 + thread_row_idx) * 128 + segment_col_idx * 16; + const void* _stmatrix_b8_ptr_3 = reinterpret_cast(reinterpret_cast(smem_p0) + stsm_offset); + uint64_t _stmatrix_b8_addr64_3; + asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(_stmatrix_b8_addr64_3) : "l"(_stmatrix_b8_ptr_3)); + uint32_t _stmatrix_b8_addr_3; + asm volatile("cvt.u32.u64 %0, %1;" : "=r"(_stmatrix_b8_addr_3) : "l"(_stmatrix_b8_addr64_3)); + asm volatile("stmatrix.sync.aligned.m16n8.x4.trans.shared.b8 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_b8_addr_3), "r"(regs_p[0]), "r"(regs_p[1]), "r"(regs_p[2]), "r"(regs_p[3]) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (tile + 2 >= num_steps_s0 * 2) { + mbarrier_wait(stats_empty0_addr, _phase_stats_empty0_0); + _phase_stats_empty0_0 ^= 1; + } + mbarrier_arrive(s_empty0_addr); + float acc_scale[4]; + #pragma unroll + for (int stat_2 = 0; stat_2 < 4; stat_2++) { + float max_delta = old_max[stat_2] - new_max[stat_2]; + acc_scale[stat_2] = 1.0f; + if (max_delta != 0.0f) { + float _exp2_13 = approx_exp2(max_delta * bmm1_scale_log2); + acc_scale[stat_2] = _exp2_13; + } + } + float2 _f2_12 = make_float2(row_sum[0], row_sum[1]); + float2 _f2_13 = make_float2(acc_scale[0], acc_scale[1]); + float2 _f2_14 = make_float2(probabilities[0], probabilities[1]); + float2 _f2_15 = make_float2(probabilities[2], probabilities[3]); + float2 _f2_16 = make_float2(probabilities[8], probabilities[9]); + float2 _f2_17 = make_float2(probabilities[10], probabilities[11]); + float2 next_sum01 = fma_f32x2_rn_ftz(_f2_12, _f2_13, _f2_14); + next_sum01 = add_f32x2(next_sum01, _f2_15); + next_sum01 = add_f32x2(next_sum01, _f2_16); + next_sum01 = add_f32x2(next_sum01, _f2_17); + row_sum[0] = next_sum01.x; + row_sum[1] = next_sum01.y; + float2 _f2_18 = make_float2(row_sum[2], row_sum[3]); + float2 _f2_19 = make_float2(acc_scale[2], acc_scale[3]); + float2 _f2_20 = make_float2(probabilities[4], probabilities[5]); + float2 _f2_21 = make_float2(probabilities[6], probabilities[7]); + float2 _f2_22 = make_float2(probabilities[12], probabilities[13]); + float2 _f2_23 = make_float2(probabilities[14], probabilities[15]); + float2 next_sum23 = fma_f32x2_rn_ftz(_f2_18, _f2_19, _f2_20); + next_sum23 = add_f32x2(next_sum23, _f2_21); + next_sum23 = add_f32x2(next_sum23, _f2_22); + next_sum23 = add_f32x2(next_sum23, _f2_23); + row_sum[2] = next_sum23.x; + row_sum[3] = next_sum23.y; + if (tile + 2 >= num_steps_s0 * 2) { + float final_stats[8]; + #pragma unroll + for (int stat_3 = 0; stat_3 < 4; stat_3++) { + final_stats[stat_3] = row_sum[stat_3]; + final_stats[stat_3 + 4] = row_max[stat_3]; + } + tmem_st_x8_f32(taddr + 32, final_stats); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(scale_ready0_addr); + } + } + } + // ---- Role: softmax1 ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 176;"); + { // softmax1_main + int cta_q_s1 = blockIdx.x / 2; + int cta_kv_s1 = blockIdx.x % 2; + int batch_s1 = blockIdx.z; + int seq_len_s1 = seq_lens[batch_s1] - (q_len - 1 - cta_q_s1); + int num_ctas_s1 = 2; + int num_steps_s1 = 2; + const int warp_in_group_1 = warp % 4; + const int wg_tid_1 = (unsigned int)(warp_in_group_1 * 32) + lane; + const int col_group_1 = wg_tid_1 % 4; + const int col_group_base_1 = col_group_1 * 4; + float row_max_1[4]; + float row_sum_1[4]; + row_max_1[0] = -3.4028235e+38f; + row_max_1[1] = -3.4028235e+38f; + row_max_1[2] = -3.4028235e+38f; + row_max_1[3] = -3.4028235e+38f; + row_sum_1[0] = 0.0f; + row_sum_1[1] = 0.0f; + row_sum_1[2] = 0.0f; + row_sum_1[3] = 0.0f; + uint32_t _amf_u_0 = __float_as_uint(-3.4028235e+38f); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_2 = _amf_u_0 ^ _amf_mask_0; + if (wg_tid_1 < 16) { + exchange1_u32[wg_tid_1] = _amf_enc_2; + } + asm volatile("barrier.sync 9, 128;" ::: "memory"); + int order_phase_1 = 0; + int order_stage_1 = 0; + unsigned int _phase_s_full1_0 = 0; + unsigned int _phase_stats_empty1_0 = 1; + #pragma unroll 1 + for (int tile_1 = 1; tile_1 < num_steps_s1 * 2; tile_1 += 2) { + mbarrier_wait(s_full1_addr, _phase_s_full1_0); + _phase_s_full1_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float scores_1[16]; + float scores_lo_1[8]; + float scores_hi_1[8]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&scores_lo_1[0])), "=r"(*reinterpret_cast(&scores_lo_1[1])), "=r"(*reinterpret_cast(&scores_lo_1[2])), "=r"(*reinterpret_cast(&scores_lo_1[3])), "=r"(*reinterpret_cast(&scores_lo_1[4])), "=r"(*reinterpret_cast(&scores_lo_1[5])), "=r"(*reinterpret_cast(&scores_lo_1[6])), "=r"(*reinterpret_cast(&scores_lo_1[7])) + : "r"(taddr + 16)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&scores_hi_1[0])), "=r"(*reinterpret_cast(&scores_hi_1[1])), "=r"(*reinterpret_cast(&scores_hi_1[2])), "=r"(*reinterpret_cast(&scores_hi_1[3])), "=r"(*reinterpret_cast(&scores_hi_1[4])), "=r"(*reinterpret_cast(&scores_hi_1[5])), "=r"(*reinterpret_cast(&scores_hi_1[6])), "=r"(*reinterpret_cast(&scores_hi_1[7])) + : "r"(taddr + 16 + 1048576)); + #pragma unroll + for (int value_idx_2 = 0; value_idx_2 < 8; value_idx_2++) { + scores_1[value_idx_2] = scores_lo_1[value_idx_2]; + scores_1[value_idx_2 + 8] = scores_hi_1[value_idx_2]; + } + int tile_token_base_1 = (cta_kv_s1 * (num_steps_s1 * 2) + tile_1) * 128; + int local_token_base_1 = (unsigned int)(warp_in_group_1 * 32) + lane / 4; + if (seq_len_s1 < tile_token_base_1 + 128) { + if (seq_len_s1 <= tile_token_base_1 + local_token_base_1) { + scores_1[0] = -3.4028235e+38f; + scores_1[1] = -3.4028235e+38f; + scores_1[4] = -3.4028235e+38f; + scores_1[5] = -3.4028235e+38f; + } + if (seq_len_s1 <= tile_token_base_1 + local_token_base_1 + 8) { + scores_1[2] = -3.4028235e+38f; + scores_1[3] = -3.4028235e+38f; + scores_1[6] = -3.4028235e+38f; + scores_1[7] = -3.4028235e+38f; + } + if (seq_len_s1 <= tile_token_base_1 + local_token_base_1 + 16) { + scores_1[8] = -3.4028235e+38f; + scores_1[9] = -3.4028235e+38f; + scores_1[12] = -3.4028235e+38f; + scores_1[13] = -3.4028235e+38f; + } + if (seq_len_s1 <= tile_token_base_1 + local_token_base_1 + 24) { + scores_1[10] = -3.4028235e+38f; + scores_1[11] = -3.4028235e+38f; + scores_1[14] = -3.4028235e+38f; + scores_1[15] = -3.4028235e+38f; + } + } + float tile_max_1[4]; + float _max_15 = max_noftz(scores_1[0], scores_1[2]); + float _max_16 = max_noftz(scores_1[8], scores_1[10]); + float _max_17 = max_noftz(_max_15, _max_16); + tile_max_1[0] = _max_17; + float _max_18 = max_noftz(scores_1[1], scores_1[3]); + float _max_19 = max_noftz(scores_1[9], scores_1[11]); + float _max_20 = max_noftz(_max_18, _max_19); + tile_max_1[1] = _max_20; + float _max_21 = max_noftz(scores_1[4], scores_1[6]); + float _max_22 = max_noftz(scores_1[12], scores_1[14]); + float _max_23 = max_noftz(_max_21, _max_22); + tile_max_1[2] = _max_23; + float _max_24 = max_noftz(scores_1[5], scores_1[7]); + float _max_25 = max_noftz(scores_1[13], scores_1[15]); + float _max_26 = max_noftz(_max_24, _max_25); + tile_max_1[3] = _max_26; + const int local_row_idx_1 = lane / 4 % 4; + float pair01_left_1 = ((local_row_idx_1 % 2 == 0) ? tile_max_1[1] : tile_max_1[0]); + float pair01_right_1 = ((local_row_idx_1 % 2 == 0) ? tile_max_1[0] : tile_max_1[1]); + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, pair01_left_1, 4); + float _max_27 = max_noftz(_shfl_xor_3, pair01_right_1); + float pair01_1 = _max_27; + float pair23_left_1 = ((local_row_idx_1 % 2 == 0) ? tile_max_1[3] : tile_max_1[2]); + float pair23_right_1 = ((local_row_idx_1 % 2 == 0) ? tile_max_1[2] : tile_max_1[3]); + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, pair23_left_1, 4); + float _max_28 = max_noftz(_shfl_xor_4, pair23_right_1); + float pair23_1 = _max_28; + float group_left_1 = ((local_row_idx_1 < 2) ? pair23_1 : pair01_1); + float group_right_1 = ((local_row_idx_1 < 2) ? pair01_1 : pair23_1); + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, group_left_1, 8); + float _max_29 = max_noftz(_shfl_xor_5, group_right_1); + float reduced_max_1 = _max_29; + uint32_t _amf_u_1 = __float_as_uint(reduced_max_1); + uint32_t _amf_mask_1 = -int32_t(_amf_u_1 >> 31) | 0x80000000u; + unsigned int _amf_enc_3 = _amf_u_1 ^ _amf_mask_1; + atomicMax(&exchange1_u32[col_group_base_1 + local_row_idx_1], _amf_enc_3); + asm volatile("barrier.sync 9, 128;" ::: "memory"); + float old_max_1[4]; + float new_max_1[4]; + #pragma unroll + for (int stat_4 = 0; stat_4 < 4; stat_4++) { + old_max_1[stat_4] = row_max_1[stat_4]; + uint32_t _amf_u_2 = exchange1_u32[col_group_base_1 + stat_4]; + uint32_t _amf_mask_2 = ((_amf_u_2 >> 31) - 1u) | 0x80000000u; + float _amf_dec_1 = __uint_as_float(_amf_u_2 ^ _amf_mask_2); + new_max_1[stat_4] = _amf_dec_1; + row_max_1[stat_4] = new_max_1[stat_4]; + } + mbarrier_wait(stats_empty1_addr, _phase_stats_empty1_0); + _phase_stats_empty1_0 ^= 1; + float local_stats_1[8]; + #pragma unroll + for (int stat_5 = 0; stat_5 < 4; stat_5++) { + local_stats_1[stat_5] = old_max_1[stat_5]; + local_stats_1[stat_5 + 4] = new_max_1[stat_5]; + } + tmem_st_x8_f32(taddr + 64, local_stats_1); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(scale_ready1_addr); + float safe_max0_1 = ((new_max_1[0] == -3.4028235e+38f) ? 0.0f : new_max_1[0]); + float safe_max1_1 = ((new_max_1[1] == -3.4028235e+38f) ? 0.0f : new_max_1[1]); + float safe_max2_1 = ((new_max_1[2] == -3.4028235e+38f) ? 0.0f : new_max_1[2]); + float safe_max3_1 = ((new_max_1[3] == -3.4028235e+38f) ? 0.0f : new_max_1[3]); + float2 _f2_24 = make_float2(bmm1_scale_log2, bmm1_scale_log2); + float2 _f2_25 = make_float2(-bmm1_scale_log2, -bmm1_scale_log2); + float2 _f2_26 = make_float2(8.8073549f, 8.8073549f); + float2 _f2_27 = make_float2(safe_max0_1, safe_max1_1); + float2 _f2_28 = make_float2(safe_max2_1, safe_max3_1); + float2 neg_max01_1 = fma_f32x2_rn_ftz(_f2_27, _f2_25, _f2_26); + float2 neg_max23_1 = fma_f32x2_rn_ftz(_f2_28, _f2_25, _f2_26); + float probabilities_1[16]; + float _fma_2 = __fmaf_rn(bmm1_scale_log2, scores_1[0], neg_max01_1.x); + probabilities_1[0] = _fma_2; + float _fma_3 = __fmaf_rn(bmm1_scale_log2, scores_1[1], neg_max01_1.y); + probabilities_1[1] = _fma_3; + float2 _f2_29 = make_float2(scores_1[2], scores_1[3]); + float2 _mul_f32x2_7; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_7) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_29)); + float2 scaled23_1 = _mul_f32x2_7; + float2 prob23_1 = add_f32x2(scaled23_1, neg_max01_1); + probabilities_1[2] = prob23_1.x; + probabilities_1[3] = prob23_1.y; + mbarrier_wait(order_p01_1_addr, order_phase_1); + float _exp2_14 = approx_exp2(probabilities_1[0]); + probabilities_1[0] = _exp2_14; + float2 _f2_30 = make_float2(scores_1[4], scores_1[5]); + float2 _mul_f32x2_8; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_8) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_30)); + float2 scaled45_1 = _mul_f32x2_8; + float2 prob45_1 = add_f32x2(scaled45_1, neg_max23_1); + probabilities_1[4] = prob45_1.x; + probabilities_1[5] = prob45_1.y; + float _exp2_15 = approx_exp2(probabilities_1[1]); + probabilities_1[1] = _exp2_15; + float _exp2_16 = approx_exp2(probabilities_1[2]); + probabilities_1[2] = _exp2_16; + float2 _f2_31 = make_float2(scores_1[6], scores_1[7]); + float2 _mul_f32x2_9; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_9) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_31)); + float2 scaled67_1 = _mul_f32x2_9; + float2 prob67_1 = add_f32x2(scaled67_1, neg_max23_1); + probabilities_1[6] = prob67_1.x; + probabilities_1[7] = prob67_1.y; + float _exp2_17 = approx_exp2(probabilities_1[3]); + probabilities_1[3] = _exp2_17; + float _exp2_18 = approx_exp2(probabilities_1[4]); + probabilities_1[4] = _exp2_18; + float2 _f2_32 = make_float2(scores_1[8], scores_1[9]); + float2 _mul_f32x2_10; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_10) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_32)); + float2 scaled89_1 = _mul_f32x2_10; + float2 prob89_1 = add_f32x2(scaled89_1, neg_max01_1); + probabilities_1[8] = prob89_1.x; + probabilities_1[9] = prob89_1.y; + float _exp2_19 = approx_exp2(probabilities_1[5]); + probabilities_1[5] = _exp2_19; + float _exp2_20 = approx_exp2(probabilities_1[6]); + probabilities_1[6] = _exp2_20; + float2 _f2_33 = make_float2(scores_1[10], scores_1[11]); + float2 _mul_f32x2_11; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_11) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_33)); + float2 scaled1011_1 = _mul_f32x2_11; + float2 prob1011_1 = add_f32x2(scaled1011_1, neg_max01_1); + probabilities_1[10] = prob1011_1.x; + probabilities_1[11] = prob1011_1.y; + float _exp2_21 = approx_exp2(probabilities_1[7]); + probabilities_1[7] = _exp2_21; + float _exp2_22 = approx_exp2(probabilities_1[8]); + probabilities_1[8] = _exp2_22; + float2 _f2_34 = make_float2(scores_1[12], scores_1[13]); + float2 _mul_f32x2_12; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_12) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_34)); + float2 scaled1213_1 = _mul_f32x2_12; + float2 prob1213_1 = add_f32x2(scaled1213_1, neg_max23_1); + probabilities_1[12] = prob1213_1.x; + probabilities_1[13] = prob1213_1.y; + float _exp2_23 = approx_exp2(probabilities_1[9]); + probabilities_1[9] = _exp2_23; + mbarrier_arrive(order_p01_0_addr); + order_phase_1 ^= 1; + float _exp2_24 = approx_exp2(probabilities_1[10]); + probabilities_1[10] = _exp2_24; + float2 _f2_35 = make_float2(scores_1[14], scores_1[15]); + float2 _mul_f32x2_13; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_13) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_35)); + float2 scaled1415_1 = _mul_f32x2_13; + float2 prob1415_1 = add_f32x2(scaled1415_1, neg_max23_1); + probabilities_1[14] = prob1415_1.x; + probabilities_1[15] = prob1415_1.y; + float _exp2_25 = approx_exp2(probabilities_1[11]); + probabilities_1[11] = _exp2_25; + unsigned int regs_p_1[4]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities_1[0]), "f"(probabilities_1[1]), + "f"(probabilities_1[2]), "f"(probabilities_1[3])); + regs_p_1[0] = _packed; + } + #pragma unroll + for (int value_idx_3 = 12; value_idx_3 < 16; value_idx_3++) { + float _exp2_26 = approx_exp2(probabilities_1[value_idx_3]); + probabilities_1[value_idx_3] = _exp2_26; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities_1[4]), "f"(probabilities_1[5]), + "f"(probabilities_1[6]), "f"(probabilities_1[7])); + regs_p_1[1] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities_1[8]), "f"(probabilities_1[9]), + "f"(probabilities_1[10]), "f"(probabilities_1[11])); + regs_p_1[2] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities_1[12]), "f"(probabilities_1[13]), + "f"(probabilities_1[14]), "f"(probabilities_1[15])); + regs_p_1[3] = _packed; + } + const int matrix_idx_1 = lane / 8; + const int matrix_row_idx_1 = matrix_idx_1 % 2; + const int matrix_col_idx_1 = matrix_idx_1 / 2; + const int thread_row_idx_1 = lane % 8; + const int segment_col_idx_1 = warp_in_group_1 * 2 + matrix_col_idx_1 ^ thread_row_idx_1; + const int stsm_offset_1 = (matrix_row_idx_1 * 8 + thread_row_idx_1) * 128 + segment_col_idx_1 * 16; + const void* _stmatrix_b8_ptr_3 = reinterpret_cast(reinterpret_cast(smem_p1) + stsm_offset_1); + uint64_t _stmatrix_b8_addr64_3; + asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(_stmatrix_b8_addr64_3) : "l"(_stmatrix_b8_ptr_3)); + uint32_t _stmatrix_b8_addr_3; + asm volatile("cvt.u32.u64 %0, %1;" : "=r"(_stmatrix_b8_addr_3) : "l"(_stmatrix_b8_addr64_3)); + asm volatile("stmatrix.sync.aligned.m16n8.x4.trans.shared.b8 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_b8_addr_3), "r"(regs_p_1[0]), "r"(regs_p_1[1]), "r"(regs_p_1[2]), "r"(regs_p_1[3]) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (tile_1 + 2 >= num_steps_s1 * 2) { + mbarrier_wait(stats_empty1_addr, _phase_stats_empty1_0); + _phase_stats_empty1_0 ^= 1; + } + mbarrier_arrive(s_empty1_addr); + float acc_scale_1[4]; + #pragma unroll + for (int stat_6 = 0; stat_6 < 4; stat_6++) { + float max_delta_1 = old_max_1[stat_6] - new_max_1[stat_6]; + acc_scale_1[stat_6] = 1.0f; + if (max_delta_1 != 0.0f) { + float _exp2_27 = approx_exp2(max_delta_1 * bmm1_scale_log2); + acc_scale_1[stat_6] = _exp2_27; + } + } + float2 _f2_36 = make_float2(row_sum_1[0], row_sum_1[1]); + float2 _f2_37 = make_float2(acc_scale_1[0], acc_scale_1[1]); + float2 _f2_38 = make_float2(probabilities_1[0], probabilities_1[1]); + float2 _f2_39 = make_float2(probabilities_1[2], probabilities_1[3]); + float2 _f2_40 = make_float2(probabilities_1[8], probabilities_1[9]); + float2 _f2_41 = make_float2(probabilities_1[10], probabilities_1[11]); + float2 next_sum01_1 = fma_f32x2_rn_ftz(_f2_36, _f2_37, _f2_38); + next_sum01_1 = add_f32x2(next_sum01_1, _f2_39); + next_sum01_1 = add_f32x2(next_sum01_1, _f2_40); + next_sum01_1 = add_f32x2(next_sum01_1, _f2_41); + row_sum_1[0] = next_sum01_1.x; + row_sum_1[1] = next_sum01_1.y; + float2 _f2_42 = make_float2(row_sum_1[2], row_sum_1[3]); + float2 _f2_43 = make_float2(acc_scale_1[2], acc_scale_1[3]); + float2 _f2_44 = make_float2(probabilities_1[4], probabilities_1[5]); + float2 _f2_45 = make_float2(probabilities_1[6], probabilities_1[7]); + float2 _f2_46 = make_float2(probabilities_1[12], probabilities_1[13]); + float2 _f2_47 = make_float2(probabilities_1[14], probabilities_1[15]); + float2 next_sum23_1 = fma_f32x2_rn_ftz(_f2_42, _f2_43, _f2_44); + next_sum23_1 = add_f32x2(next_sum23_1, _f2_45); + next_sum23_1 = add_f32x2(next_sum23_1, _f2_46); + next_sum23_1 = add_f32x2(next_sum23_1, _f2_47); + row_sum_1[2] = next_sum23_1.x; + row_sum_1[3] = next_sum23_1.y; + if (tile_1 + 2 >= num_steps_s1 * 2) { + float final_stats_1[8]; + #pragma unroll + for (int stat_7 = 0; stat_7 < 4; stat_7++) { + final_stats_1[stat_7] = row_sum_1[stat_7]; + final_stats_1[stat_7 + 4] = row_max_1[stat_7]; + } + tmem_st_x8_f32(taddr + 64, final_stats_1); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(scale_ready1_addr); + } + } + } + // ---- Role: correction ---- + } else if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + { // correction_main + int cta_q_c = blockIdx.x / 2; + int cta_kv_c = blockIdx.x % 2; + int y_c = blockIdx.y; + int head_group_c = y_c / 4; + int hv_partition_c = y_c % 4; + int batch_c = blockIdx.z; + int seq_len_c = seq_lens[batch_c] - (q_len - 1 - cta_q_c); + int num_ctas_c = 2; + int num_steps_c = 2; + int num_tiles_c = num_steps_c * 2; + const int warp_c = warp % 4; + const int corr_tid = (unsigned int)(warp_c * 32) + lane; + const int tmem_row_c = warp_c * 32 << 16; + const int corr_head_base_c = corr_tid % 4 * 4; + unsigned int _phase_scale_ready0_0 = 0; + unsigned int _phase_o_full0_0 = 0; + unsigned int _phase_scale_ready1_0 = 0; + unsigned int _phase_o_full1_0 = 0; + #pragma unroll 1 + for (int tile_c = 0; tile_c < num_tiles_c; tile_c++) { + if (tile_c % 2 == 0) { + mbarrier_wait(scale_ready0_addr, _phase_scale_ready0_0); + _phase_scale_ready0_0 ^= 1; + float corr_scale0[4]; + if (tile_c > 0) { + float _tmem_load_0[8]; + tmem_ld_x8(&_tmem_load_0[0], taddr + 32); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int head_local_c = 0; head_local_c < 4; head_local_c++) { + float corr_delta0_c = _tmem_load_0[head_local_c] - _tmem_load_0[head_local_c + 4]; + float _exp2_28 = approx_exp2(corr_delta0_c * bmm1_scale_log2); + corr_scale0[head_local_c] = ((corr_delta0_c != 0.0f) ? _exp2_28 : 1.0f); + } + } + mbarrier_arrive(stats_empty0_addr); + if (tile_c > 0) { + mbarrier_wait(o_full0_addr, _phase_o_full0_0); + _phase_o_full0_0 ^= 1; + float previous0_lo[8]; + float previous0_hi[8]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&previous0_lo[0])), "=r"(*reinterpret_cast(&previous0_lo[1])), "=r"(*reinterpret_cast(&previous0_lo[2])), "=r"(*reinterpret_cast(&previous0_lo[3])), "=r"(*reinterpret_cast(&previous0_lo[4])), "=r"(*reinterpret_cast(&previous0_lo[5])), "=r"(*reinterpret_cast(&previous0_lo[6])), "=r"(*reinterpret_cast(&previous0_lo[7])) + : "r"(taddr + 96)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&previous0_hi[0])), "=r"(*reinterpret_cast(&previous0_hi[1])), "=r"(*reinterpret_cast(&previous0_hi[2])), "=r"(*reinterpret_cast(&previous0_hi[3])), "=r"(*reinterpret_cast(&previous0_hi[4])), "=r"(*reinterpret_cast(&previous0_hi[5])), "=r"(*reinterpret_cast(&previous0_hi[6])), "=r"(*reinterpret_cast(&previous0_hi[7])) + : "r"(taddr + 96 + 1048576)); + float2 _f2_48 = make_float2(corr_scale0[0], corr_scale0[1]); + float2 _f2_49 = make_float2(corr_scale0[2], corr_scale0[3]); + int corr_pred0 = 0; + #pragma unroll + for (int head_local_c_1 = 0; head_local_c_1 < 4; head_local_c_1++) { + if (corr_scale0[head_local_c_1] != 1.0f) { + corr_pred0 = 1; + } + } + int _vote_0 = __any_sync(0xFFFFFFFF, corr_pred0 != 0); + if (_vote_0 != 0) { + #pragma unroll + for (int reg_base_c = 0; reg_base_c < 8; reg_base_c += 4) { + float2 _f2_50 = make_float2(previous0_lo[reg_base_c], previous0_lo[reg_base_c + 1]); + float2 _f2_51 = make_float2(previous0_lo[reg_base_c + 2], previous0_lo[reg_base_c + 3]); + float2 _f2_52 = make_float2(previous0_hi[reg_base_c], previous0_hi[reg_base_c + 1]); + float2 _f2_53 = make_float2(previous0_hi[reg_base_c + 2], previous0_hi[reg_base_c + 3]); + float2 _f32x2_mul_b_0 = ((reg_base_c == 0) ? _f2_48 : _f2_49); + float2 _mul_f32x2_14; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_14) : "l"(*(const unsigned long long*)&_f2_50), "l"(*(const unsigned long long*)&_f32x2_mul_b_0)); + float2 prev0_lo01_scaled_f2 = _mul_f32x2_14; + float2 _f32x2_mul_b_1 = ((reg_base_c == 0) ? _f2_48 : _f2_49); + float2 _mul_f32x2_15; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_15) : "l"(*(const unsigned long long*)&_f2_51), "l"(*(const unsigned long long*)&_f32x2_mul_b_1)); + float2 prev0_lo23_scaled_f2 = _mul_f32x2_15; + float2 _f32x2_mul_b_2 = ((reg_base_c == 0) ? _f2_48 : _f2_49); + float2 _mul_f32x2_16; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_16) : "l"(*(const unsigned long long*)&_f2_52), "l"(*(const unsigned long long*)&_f32x2_mul_b_2)); + float2 prev0_hi01_scaled_f2 = _mul_f32x2_16; + float2 _f32x2_mul_b_3 = ((reg_base_c == 0) ? _f2_48 : _f2_49); + float2 _mul_f32x2_17; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_17) : "l"(*(const unsigned long long*)&_f2_53), "l"(*(const unsigned long long*)&_f32x2_mul_b_3)); + float2 prev0_hi23_scaled_f2 = _mul_f32x2_17; + previous0_lo[reg_base_c] = prev0_lo01_scaled_f2.x; + previous0_lo[reg_base_c + 1] = prev0_lo01_scaled_f2.y; + previous0_lo[reg_base_c + 2] = prev0_lo23_scaled_f2.x; + previous0_lo[reg_base_c + 3] = prev0_lo23_scaled_f2.y; + previous0_hi[reg_base_c] = prev0_hi01_scaled_f2.x; + previous0_hi[reg_base_c + 1] = prev0_hi01_scaled_f2.y; + previous0_hi[reg_base_c + 2] = prev0_hi23_scaled_f2.x; + previous0_hi[reg_base_c + 3] = prev0_hi23_scaled_f2.y; + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x2.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(taddr + 96), "r"(*reinterpret_cast(&previous0_lo[0])), "r"(*reinterpret_cast(&previous0_lo[1])), "r"(*reinterpret_cast(&previous0_lo[2])), "r"(*reinterpret_cast(&previous0_lo[3])), "r"(*reinterpret_cast(&previous0_lo[4])), "r"(*reinterpret_cast(&previous0_lo[5])), "r"(*reinterpret_cast(&previous0_lo[6])), "r"(*reinterpret_cast(&previous0_lo[7]))); + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x2.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(taddr + 96 + 1048576), "r"(*reinterpret_cast(&previous0_hi[0])), "r"(*reinterpret_cast(&previous0_hi[1])), "r"(*reinterpret_cast(&previous0_hi[2])), "r"(*reinterpret_cast(&previous0_hi[3])), "r"(*reinterpret_cast(&previous0_hi[4])), "r"(*reinterpret_cast(&previous0_hi[5])), "r"(*reinterpret_cast(&previous0_hi[6])), "r"(*reinterpret_cast(&previous0_hi[7]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + mbarrier_arrive(o_empty0_addr); + } + } else { + mbarrier_wait(scale_ready1_addr, _phase_scale_ready1_0); + _phase_scale_ready1_0 ^= 1; + float corr_scale1[4]; + if (tile_c > 1) { + float _tmem_load_1[8]; + tmem_ld_x8(&_tmem_load_1[0], taddr + 64); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int head_local_c_2 = 0; head_local_c_2 < 4; head_local_c_2++) { + float corr_delta1_c = _tmem_load_1[head_local_c_2] - _tmem_load_1[head_local_c_2 + 4]; + float _exp2_29 = approx_exp2(corr_delta1_c * bmm1_scale_log2); + corr_scale1[head_local_c_2] = ((corr_delta1_c != 0.0f) ? _exp2_29 : 1.0f); + } + } + mbarrier_arrive(stats_empty1_addr); + if (tile_c > 1) { + mbarrier_wait(o_full1_addr, _phase_o_full1_0); + _phase_o_full1_0 ^= 1; + float previous1_lo[8]; + float previous1_hi[8]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&previous1_lo[0])), "=r"(*reinterpret_cast(&previous1_lo[1])), "=r"(*reinterpret_cast(&previous1_lo[2])), "=r"(*reinterpret_cast(&previous1_lo[3])), "=r"(*reinterpret_cast(&previous1_lo[4])), "=r"(*reinterpret_cast(&previous1_lo[5])), "=r"(*reinterpret_cast(&previous1_lo[6])), "=r"(*reinterpret_cast(&previous1_lo[7])) + : "r"(taddr + 112)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&previous1_hi[0])), "=r"(*reinterpret_cast(&previous1_hi[1])), "=r"(*reinterpret_cast(&previous1_hi[2])), "=r"(*reinterpret_cast(&previous1_hi[3])), "=r"(*reinterpret_cast(&previous1_hi[4])), "=r"(*reinterpret_cast(&previous1_hi[5])), "=r"(*reinterpret_cast(&previous1_hi[6])), "=r"(*reinterpret_cast(&previous1_hi[7])) + : "r"(taddr + 112 + 1048576)); + float2 _f2_54 = make_float2(corr_scale1[0], corr_scale1[1]); + float2 _f2_55 = make_float2(corr_scale1[2], corr_scale1[3]); + int corr_pred1 = 0; + #pragma unroll + for (int head_local_c_3 = 0; head_local_c_3 < 4; head_local_c_3++) { + if (corr_scale1[head_local_c_3] != 1.0f) { + corr_pred1 = 1; + } + } + int _vote_1 = __any_sync(0xFFFFFFFF, corr_pred1 != 0); + if (_vote_1 != 0) { + #pragma unroll + for (int reg_base_c_1 = 0; reg_base_c_1 < 8; reg_base_c_1 += 4) { + float2 _f2_56 = make_float2(previous1_lo[reg_base_c_1], previous1_lo[reg_base_c_1 + 1]); + float2 _f2_57 = make_float2(previous1_lo[reg_base_c_1 + 2], previous1_lo[reg_base_c_1 + 3]); + float2 _f2_58 = make_float2(previous1_hi[reg_base_c_1], previous1_hi[reg_base_c_1 + 1]); + float2 _f2_59 = make_float2(previous1_hi[reg_base_c_1 + 2], previous1_hi[reg_base_c_1 + 3]); + float2 _f32x2_mul_b_4 = ((reg_base_c_1 == 0) ? _f2_54 : _f2_55); + float2 _mul_f32x2_18; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_18) : "l"(*(const unsigned long long*)&_f2_56), "l"(*(const unsigned long long*)&_f32x2_mul_b_4)); + float2 prev1_lo01_scaled_f2 = _mul_f32x2_18; + float2 _f32x2_mul_b_5 = ((reg_base_c_1 == 0) ? _f2_54 : _f2_55); + float2 _mul_f32x2_19; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_19) : "l"(*(const unsigned long long*)&_f2_57), "l"(*(const unsigned long long*)&_f32x2_mul_b_5)); + float2 prev1_lo23_scaled_f2 = _mul_f32x2_19; + float2 _f32x2_mul_b_6 = ((reg_base_c_1 == 0) ? _f2_54 : _f2_55); + float2 _mul_f32x2_20; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_20) : "l"(*(const unsigned long long*)&_f2_58), "l"(*(const unsigned long long*)&_f32x2_mul_b_6)); + float2 prev1_hi01_scaled_f2 = _mul_f32x2_20; + float2 _f32x2_mul_b_7 = ((reg_base_c_1 == 0) ? _f2_54 : _f2_55); + float2 _mul_f32x2_21; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_21) : "l"(*(const unsigned long long*)&_f2_59), "l"(*(const unsigned long long*)&_f32x2_mul_b_7)); + float2 prev1_hi23_scaled_f2 = _mul_f32x2_21; + previous1_lo[reg_base_c_1] = prev1_lo01_scaled_f2.x; + previous1_lo[reg_base_c_1 + 1] = prev1_lo01_scaled_f2.y; + previous1_lo[reg_base_c_1 + 2] = prev1_lo23_scaled_f2.x; + previous1_lo[reg_base_c_1 + 3] = prev1_lo23_scaled_f2.y; + previous1_hi[reg_base_c_1] = prev1_hi01_scaled_f2.x; + previous1_hi[reg_base_c_1 + 1] = prev1_hi01_scaled_f2.y; + previous1_hi[reg_base_c_1 + 2] = prev1_hi23_scaled_f2.x; + previous1_hi[reg_base_c_1 + 3] = prev1_hi23_scaled_f2.y; + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x2.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(taddr + 112), "r"(*reinterpret_cast(&previous1_lo[0])), "r"(*reinterpret_cast(&previous1_lo[1])), "r"(*reinterpret_cast(&previous1_lo[2])), "r"(*reinterpret_cast(&previous1_lo[3])), "r"(*reinterpret_cast(&previous1_lo[4])), "r"(*reinterpret_cast(&previous1_lo[5])), "r"(*reinterpret_cast(&previous1_lo[6])), "r"(*reinterpret_cast(&previous1_lo[7]))); + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x2.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(taddr + 112 + 1048576), "r"(*reinterpret_cast(&previous1_hi[0])), "r"(*reinterpret_cast(&previous1_hi[1])), "r"(*reinterpret_cast(&previous1_hi[2])), "r"(*reinterpret_cast(&previous1_hi[3])), "r"(*reinterpret_cast(&previous1_hi[4])), "r"(*reinterpret_cast(&previous1_hi[5])), "r"(*reinterpret_cast(&previous1_hi[6])), "r"(*reinterpret_cast(&previous1_hi[7]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + mbarrier_arrive(o_empty1_addr); + } + } + } + mbarrier_wait(scale_ready0_addr, _phase_scale_ready0_0); + _phase_scale_ready0_0 ^= 1; + float _tmem_load_2[8]; + tmem_ld_x8(&_tmem_load_2[0], taddr + 32); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + mbarrier_arrive(stats_empty0_addr); + mbarrier_wait(scale_ready1_addr, _phase_scale_ready1_0); + _phase_scale_ready1_0 ^= 1; + float _tmem_load_3[8]; + tmem_ld_x8(&_tmem_load_3[0], taddr + 64); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + mbarrier_arrive(stats_empty1_addr); + mbarrier_wait(o_full0_addr, _phase_o_full0_0); + _phase_o_full0_0 ^= 1; + if (num_tiles_c > 1) { + mbarrier_wait(o_full1_addr, _phase_o_full1_0); + _phase_o_full1_0 ^= 1; + } + asm volatile("tcgen05.fence::after_thread_sync;"); + float out0_lo[8]; + float out0_hi[8]; + float out1_lo[8]; + float out1_hi[8]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&out0_lo[0])), "=r"(*reinterpret_cast(&out0_lo[1])), "=r"(*reinterpret_cast(&out0_lo[2])), "=r"(*reinterpret_cast(&out0_lo[3])), "=r"(*reinterpret_cast(&out0_lo[4])), "=r"(*reinterpret_cast(&out0_lo[5])), "=r"(*reinterpret_cast(&out0_lo[6])), "=r"(*reinterpret_cast(&out0_lo[7])) + : "r"(taddr + 96)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&out0_hi[0])), "=r"(*reinterpret_cast(&out0_hi[1])), "=r"(*reinterpret_cast(&out0_hi[2])), "=r"(*reinterpret_cast(&out0_hi[3])), "=r"(*reinterpret_cast(&out0_hi[4])), "=r"(*reinterpret_cast(&out0_hi[5])), "=r"(*reinterpret_cast(&out0_hi[6])), "=r"(*reinterpret_cast(&out0_hi[7])) + : "r"(taddr + 96 + 1048576)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&out1_lo[0])), "=r"(*reinterpret_cast(&out1_lo[1])), "=r"(*reinterpret_cast(&out1_lo[2])), "=r"(*reinterpret_cast(&out1_lo[3])), "=r"(*reinterpret_cast(&out1_lo[4])), "=r"(*reinterpret_cast(&out1_lo[5])), "=r"(*reinterpret_cast(&out1_lo[6])), "=r"(*reinterpret_cast(&out1_lo[7])) + : "r"(taddr + 112)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&out1_hi[0])), "=r"(*reinterpret_cast(&out1_hi[1])), "=r"(*reinterpret_cast(&out1_hi[2])), "=r"(*reinterpret_cast(&out1_hi[3])), "=r"(*reinterpret_cast(&out1_hi[4])), "=r"(*reinterpret_cast(&out1_hi[5])), "=r"(*reinterpret_cast(&out1_hi[6])), "=r"(*reinterpret_cast(&out1_hi[7])) + : "r"(taddr + 112 + 1048576)); + int logical_cta_c = ((batch_c * 32 + y_c) * max_num_ctas_q + cta_q_c) * max_num_ctas_kv + cta_kv_c; + int partial_o_base_c = logical_cta_c * 16 * 128; + int stats_base_c = logical_cta_c * 16; + float final_merged_max_c[4]; + float final_local_sum_c[4]; + float final_scale0_c[4]; + float final_scale1_c[4]; + #pragma unroll + for (int head_local_c_4 = 0; head_local_c_4 < 4; head_local_c_4++) { + float _max_30 = max_noftz(_tmem_load_2[head_local_c_4 + 4], _tmem_load_3[head_local_c_4 + 4]); + float merged_max_c = _max_30; + float _exp2_30 = approx_exp2((_tmem_load_2[head_local_c_4 + 4] - merged_max_c) * bmm1_scale_log2); + float weight0_c = ((_tmem_load_2[head_local_c_4 + 4] == -3.4028235e+38f) ? 0.0f : _exp2_30); + float _exp2_31 = approx_exp2((_tmem_load_3[head_local_c_4 + 4] - merged_max_c) * bmm1_scale_log2); + float weight1_c = ((_tmem_load_3[head_local_c_4 + 4] == -3.4028235e+38f) ? 0.0f : _exp2_31); + final_merged_max_c[head_local_c_4] = merged_max_c; + final_local_sum_c[head_local_c_4] = _tmem_load_2[head_local_c_4] * weight0_c + _tmem_load_3[head_local_c_4] * weight1_c; + final_scale0_c[head_local_c_4] = weight0_c * bmm2_scale / 448.0f; + final_scale1_c[head_local_c_4] = weight1_c * bmm2_scale / 448.0f; + } + #pragma unroll + for (int head_local_c_5 = 0; head_local_c_5 < 4; head_local_c_5++) { + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, final_local_sum_c[head_local_c_5], 16); + final_local_sum_c[head_local_c_5] = final_local_sum_c[head_local_c_5] + _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, final_local_sum_c[head_local_c_5], 8); + final_local_sum_c[head_local_c_5] = final_local_sum_c[head_local_c_5] + _shfl_xor_7; + float _shfl_xor_8 = __shfl_xor_sync(0xFFFFFFFF, final_local_sum_c[head_local_c_5], 4); + final_local_sum_c[head_local_c_5] = final_local_sum_c[head_local_c_5] + _shfl_xor_8; + if (lane < 4) { + corr_stats_reduce[warp_c * 16 + corr_head_base_c + head_local_c_5] = final_local_sum_c[head_local_c_5]; + } + } + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float final_global_sum_c[4]; + #pragma unroll + for (int head_local_c_6 = 0; head_local_c_6 < 4; head_local_c_6++) { + int head_offset_c = corr_head_base_c + head_local_c_6; + final_global_sum_c[head_local_c_6] = corr_stats_reduce[head_offset_c] + corr_stats_reduce[16 + head_offset_c] + corr_stats_reduce[32 + head_offset_c] + corr_stats_reduce[48 + head_offset_c]; + } + float merged_o_c[16]; + #pragma unroll + for (int reg_base_c_2 = 0; reg_base_c_2 < 8; reg_base_c_2 += 4) { + const int scale_idx_c = ((reg_base_c_2 == 0) ? 0 : 2); + float2 _f2_60 = make_float2(final_scale0_c[scale_idx_c], final_scale0_c[scale_idx_c + 1]); + float2 _f2_61 = make_float2(final_scale1_c[scale_idx_c], final_scale1_c[scale_idx_c + 1]); + #pragma unroll + for (int pair_off_c = 0; pair_off_c < 4; pair_off_c += 2) { + const int reg_c = reg_base_c_2 + pair_off_c; + float2 _f2_62 = make_float2(out0_lo[reg_c], out0_lo[reg_c + 1]); + float2 _f2_63 = make_float2(out1_lo[reg_c], out1_lo[reg_c + 1]); + float2 _f2_64 = make_float2(out0_hi[reg_c], out0_hi[reg_c + 1]); + float2 _f2_65 = make_float2(out1_hi[reg_c], out1_hi[reg_c + 1]); + float2 _mul_f32x2_22; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_22) : "l"(*(const unsigned long long*)&_f2_63), "l"(*(const unsigned long long*)&_f2_61)); + float2 out1_lo_scaled_c = _mul_f32x2_22; + float2 _mul_f32x2_23; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_23) : "l"(*(const unsigned long long*)&_f2_65), "l"(*(const unsigned long long*)&_f2_61)); + float2 out1_hi_scaled_c = _mul_f32x2_23; + float2 merged_lo_pair_c = fma_f32x2_rn_ftz(_f2_62, _f2_60, out1_lo_scaled_c); + float2 merged_hi_pair_c = fma_f32x2_rn_ftz(_f2_64, _f2_60, out1_hi_scaled_c); + merged_o_c[reg_c] = merged_lo_pair_c.x; + merged_o_c[reg_c + 1] = merged_lo_pair_c.y; + merged_o_c[reg_c + 8] = merged_hi_pair_c.x; + merged_o_c[reg_c + 9] = merged_hi_pair_c.y; + } + } + unsigned int regs_o_c[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(merged_o_c[_lp*2 + 0], merged_o_c[_lp*2+1 + 0])); + regs_o_c[_lp] = *(uint32_t*)&_bf2; + } + int slice_idx_c = warp_c / 2; + int warp_idx_in_slice_c = warp_c % 2; + int mtx_idx_c = lane / 8; + int mtx_row_idx_c = mtx_idx_c / 2; + int mtx_col_idx_c = warp_idx_in_slice_c * 4 + mtx_idx_c % 2; + int thr_row_idx_c = lane % 8; + #pragma unroll + for (int stsm_idx_c = 0; stsm_idx_c < 2; stsm_idx_c++) { + int seg_col_idx_c = mtx_col_idx_c + stsm_idx_c * 2 ^ thr_row_idx_c; + int stsm_offset_c = slice_idx_c * 16 * 128 + (mtx_row_idx_c * 8 + thr_row_idx_c) * 128 + seg_col_idx_c * 16; + const int stsm_reg_c = stsm_idx_c * 4; + const void* _stmatrix_ptr_8 = reinterpret_cast(reinterpret_cast(smem_o) + stsm_offset_c); + uint64_t _stmatrix_addr64_8; + asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(_stmatrix_addr64_8) : "l"(_stmatrix_ptr_8)); + uint32_t _stmatrix_addr_8; + asm volatile("cvt.u32.u64 %0, %1;" : "=r"(_stmatrix_addr_8) : "l"(_stmatrix_addr64_8)); + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_8), "r"(*reinterpret_cast(®s_o_c[stsm_reg_c])), "r"(*reinterpret_cast(®s_o_c[stsm_reg_c + 1])), "r"(*reinterpret_cast(®s_o_c[stsm_reg_c + 2])), "r"(*reinterpret_cast(®s_o_c[stsm_reg_c + 3])) + : "memory"); + } + asm volatile("barrier.sync 3, 128;" ::: "memory"); + #pragma unroll + for (int copy_idx_c = 0; copy_idx_c < 2; copy_idx_c++) { + int copy_base_c = corr_tid * 16 + copy_idx_c * 128 * 16; + int copy_smem_row_c = copy_base_c / 128; + int copy_dst_row_c = copy_smem_row_c % 16; + int copy_dst_col_bytes_c = copy_smem_row_c / 16 * 128 + copy_base_c % 128; + int copy_smem_offset_c = copy_base_c ^ copy_smem_row_c % 8 * 16; + unsigned int copy_vec_c[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(©_vec_c[0])), "=r"(*reinterpret_cast(©_vec_c[(0) + 1])), "=r"(*reinterpret_cast(©_vec_c[(0) + 2])), "=r"(*reinterpret_cast(©_vec_c[(0) + 3])) + : "r"(smem_o_addr + (unsigned int)copy_smem_offset_c)); + reinterpret_cast(partial_O + (partial_o_base_c + copy_dst_row_c * 128 + copy_dst_col_bytes_c / 2))[0] = reinterpret_cast(copy_vec_c)[0]; + } + if (corr_tid < 4) { + #pragma unroll + for (int head_local_c_7 = 0; head_local_c_7 < 4; head_local_c_7++) { + int head_c = 2 * (corr_tid % 4) + 8 * (head_local_c_7 / 2) + head_local_c_7 % 2; + partial_stats[(stats_base_c + head_c) * 2] = final_merged_max_c[head_local_c_7]; + partial_stats[(stats_base_c + head_c) * 2 + 1] = final_global_sum_c[head_local_c_7]; + } + } + mbarrier_arrive(o_empty0_addr); + if (num_tiles_c > 1) { + mbarrier_arrive(o_empty1_addr); + } + int reduction_group_c = (batch_c * 32 + y_c) * max_num_ctas_q + cta_q_c; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + if (corr_tid == 0) { + uint32_t _atomic_inc_old_0; + asm volatile("atom.acq_rel.gpu.global.inc.u32 %0, [%1], %2;" + : "=r"(_atomic_inc_old_0) : "l"(&completion[reduction_group_c]), "r"(static_cast(1)) : "memory"); + unsigned int old_ticket_c = _atomic_inc_old_0; + reduction_ticket[0] = 1 - (int)old_ticket_c; + if (old_ticket_c == 0) { + { + unsigned int* _global_wait_zero_p_9 = (reinterpret_cast(completion) + (reduction_group_c)); + unsigned int _global_wait_zero_v_9; + do { + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(_global_wait_zero_v_9) : "l"(_global_wait_zero_p_9) : "memory"); + } while (_global_wait_zero_v_9 != 0u); + } + } + } + asm volatile("barrier.sync 3, 128;" ::: "memory"); + int reduction_cta_c = reduction_ticket[0]; + int reduction_row_c = reduction_cta_c * 8 + corr_tid / 16; + int reduction_col_c = corr_tid * 8 % 128; + int stats_group_base_c = reduction_group_c * max_num_ctas_kv * 16 * 2; + int stats_row_c = reduction_row_c * 2; + float _vec_load_0[2]; + { + float2 _v2_10 = *reinterpret_cast(partial_stats + stats_group_base_c + stats_row_c); + _vec_load_0[0] = _v2_10.x; + _vec_load_0[0 + 1] = _v2_10.y; + } + float _vec_load_1[2]; + { + float2 _v2_11 = *reinterpret_cast(partial_stats + stats_group_base_c + 32 + stats_row_c); + _vec_load_1[0] = _v2_11.x; + _vec_load_1[0 + 1] = _v2_11.y; + } + float max0_c = _vec_load_0[0]; + float sum0_c = _vec_load_0[1]; + float max1_c = _vec_load_1[0]; + float sum1_c = _vec_load_1[1]; + float _max_31 = max_noftz(max0_c, max1_c); + float global_max_c = _max_31; + float _exp2_32 = approx_exp2((max0_c - global_max_c) * bmm1_scale_log2); + float cta_scale0_c = _exp2_32; + float _exp2_33 = approx_exp2((max1_c - global_max_c) * bmm1_scale_log2); + float cta_scale1_c = _exp2_33; + float global_sum_c = sum0_c * cta_scale0_c + sum1_c * cta_scale1_c; + float normalized_scale_c = 448.0f / global_sum_c; + int partial_group_base_c = reduction_group_c * max_num_ctas_kv * 16 * 128; + int partial_row_col_c = reduction_row_c * 128 + reduction_col_c; + float _vec_load_2[8]; + { + const uint4* _vptr_12 = reinterpret_cast(partial_O + (partial_group_base_c + partial_row_col_c) + 0); + uint4 _vld_12[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_12[_blk] = _vptr_12[_blk]; + uint32_t* _vpairs_12 = reinterpret_cast(&_vld_12[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&_vec_load_2[0 + _blk * 8 + _pair * 2])[0]), "=f"((&_vec_load_2[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_12[_pair])); + } + } + } + float _vec_load_3[8]; + { + const uint4* _vptr_13 = reinterpret_cast(partial_O + (partial_group_base_c + 2048 + partial_row_col_c) + 0); + uint4 _vld_13[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_13[_blk] = _vptr_13[_blk]; + uint32_t* _vpairs_13 = reinterpret_cast(&_vld_13[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&_vec_load_3[0 + _blk * 8 + _pair * 2])[0]), "=f"((&_vec_load_3[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_13[_pair])); + } + } + } + float reduced_c[8]; + float2 _f2_66 = make_float2(cta_scale0_c, cta_scale0_c); + float2 _f2_67 = make_float2(cta_scale1_c, cta_scale1_c); + float2 _f2_68 = make_float2(normalized_scale_c, normalized_scale_c); + #pragma unroll + for (int elem_c = 0; elem_c < 8; elem_c += 2) { + float2 _f2_69 = make_float2(_vec_load_2[elem_c], _vec_load_2[elem_c + 1]); + float2 _f2_70 = make_float2(_vec_load_3[elem_c], _vec_load_3[elem_c + 1]); + float2 _mul_f32x2_24; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_24) : "l"(*(const unsigned long long*)&_f2_70), "l"(*(const unsigned long long*)&_f2_67)); + float2 merged_pair_c = fma_f32x2_rn_ftz(_f2_69, _f2_66, _mul_f32x2_24); + float2 _mul_f32x2_25; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_25) : "l"(*(const unsigned long long*)&merged_pair_c), "l"(*(const unsigned long long*)&_f2_68)); + float2 reduced_pair_c = _mul_f32x2_25; + reduced_c[elem_c] = reduced_pair_c.x; + reduced_c[elem_c + 1] = reduced_pair_c.y; + } + int output_head_c = head_group_c * 16 + reduction_row_c; + int output_base_c = ((batch_c * q_len + cta_q_c) * 128 + output_head_c) * 512 + hv_partition_c * 128 + reduction_col_c; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(reduced_c[0 + 0], reduced_c[0 + 1]); + _pk[1] = __floats2bfloat162_rn(reduced_c[0 + 2], reduced_c[0 + 3]); + _pk[2] = __floats2bfloat162_rn(reduced_c[0 + 4], reduced_c[0 + 5]); + _pk[3] = __floats2bfloat162_rn(reduced_c[0 + 6], reduced_c[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + output_base_c))[0]) = *reinterpret_cast(&_pk[0]); + } + asm volatile("barrier.sync 13, 128;" ::: "memory"); + if (warp_c == 0) { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: mma_warp ---- + } else if (warp == 12) { + { // mma_warp_main + int cta_q_m = blockIdx.x / 2; + int batch_m = blockIdx.z; + int seq_len_m = seq_lens[batch_m] - (q_len - 1 - cta_q_m); + int num_ctas_m = 2; + int num_steps_m = 2; + int kv_stage_m = 0; + unsigned int _phase_q_full_0 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + unsigned int _phase_s_empty0_0 = 1; + mbarrier_wait(s_empty0_addr, _phase_s_empty0_0); + _phase_s_empty0_0 ^= 1; + unsigned int _phase_kv_full = 0; + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_0 = make_warp_uniform((((smem_q0_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134479888, 0); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_1 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_1 = make_warp_uniform((((smem_q1_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_2 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_2 = make_warp_uniform((((smem_q2_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_2); + uint64_t _mma_ss_b_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_2); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_3 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_3 = make_warp_uniform((((smem_q3_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_3); + uint64_t _mma_ss_b_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_3); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_4 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_4 = make_warp_uniform((((smem_q4_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_4); + uint64_t _mma_ss_b_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_4); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + elect_commit(s_full0_addr); + unsigned int _phase_s_empty1_0 = 1; + mbarrier_wait(s_empty1_addr, _phase_s_empty1_0); + _phase_s_empty1_0 ^= 1; + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_5 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_5 = make_warp_uniform((((smem_q0_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_5); + uint64_t _mma_ss_b_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_5); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134479888, 0); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_6 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_6 = make_warp_uniform((((smem_q1_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_6 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_6); + uint64_t _mma_ss_b_desc_6 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_6); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_6, _mma_ss_b_desc_6, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_6, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_6, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_6, _mma_ss_b_desc_6, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_6, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_6, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_6, _mma_ss_b_desc_6, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_6, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_6, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_6, _mma_ss_b_desc_6, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_7 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_7 = make_warp_uniform((((smem_q2_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_7 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_7); + uint64_t _mma_ss_b_desc_7 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_7); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_7, _mma_ss_b_desc_7, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_7, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_7, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_7, _mma_ss_b_desc_7, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_7, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_7, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_7, _mma_ss_b_desc_7, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_7, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_7, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_7, _mma_ss_b_desc_7, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_8 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_8 = make_warp_uniform((((smem_q3_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_8 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_8); + uint64_t _mma_ss_b_desc_8 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_8); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_8, _mma_ss_b_desc_8, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_8, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_8, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_8, _mma_ss_b_desc_8, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_8, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_8, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_8, _mma_ss_b_desc_8, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_8, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_8, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_8, _mma_ss_b_desc_8, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_9 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_9 = make_warp_uniform((((smem_q4_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_9 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_9); + uint64_t _mma_ss_b_desc_9 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_9); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_9, _mma_ss_b_desc_9, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_9, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_9, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_9, _mma_ss_b_desc_9, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_9, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_9, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_9, _mma_ss_b_desc_9, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_9, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_9, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_9, _mma_ss_b_desc_9, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + elect_commit(s_full1_addr); + unsigned int _phase_o_empty0_0 = 1; + unsigned int _phase_o_empty1_0 = 1; + #pragma unroll 1 + for (int pair_m = 0; pair_m < num_steps_m - 1; pair_m++) { + mbarrier_wait(s_empty0_addr, _phase_s_empty0_0); + _phase_s_empty0_0 ^= 1; + mbarrier_wait(o_empty0_addr, _phase_o_empty0_0); + _phase_o_empty0_0 ^= 1; + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_10 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (kv_stage_m) * 1024); + int _mma_b_lo_10 = make_warp_uniform((((smem_p0_addr) >> 4) & 0x3FFF) | 0x800000); + { + uint64_t _mma_ss_a_desc_10 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_10); + uint64_t _mma_ss_b_desc_10 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_10); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_10, _mma_ss_b_desc_10, 134512656, ((pair_m == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_10, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_10, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_10, _mma_ss_b_desc_10, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_10, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_10, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_10, _mma_ss_b_desc_10, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_10, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_10, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_10, _mma_ss_b_desc_10, 134512656, 1); + } + } + elect_commit2(kv_empty_addr + (kv_stage_m) * 8, o_full0_addr); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_11 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_11 = make_warp_uniform((((smem_q0_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_11 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_11); + uint64_t _mma_ss_b_desc_11 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_11); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_11, _mma_ss_b_desc_11, 134479888, 0); + } + incr_smem_desc_lo(_mma_ss_a_desc_11, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_11, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_11, _mma_ss_b_desc_11, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_11, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_11, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_11, _mma_ss_b_desc_11, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_11, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_11, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_11, _mma_ss_b_desc_11, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_12 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_12 = make_warp_uniform((((smem_q1_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_12 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_12); + uint64_t _mma_ss_b_desc_12 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_12); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_12, _mma_ss_b_desc_12, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_12, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_12, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_12, _mma_ss_b_desc_12, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_12, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_12, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_12, _mma_ss_b_desc_12, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_12, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_12, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_12, _mma_ss_b_desc_12, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_13 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_13 = make_warp_uniform((((smem_q2_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_13 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_13); + uint64_t _mma_ss_b_desc_13 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_13); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_13, _mma_ss_b_desc_13, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_13, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_13, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_13, _mma_ss_b_desc_13, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_13, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_13, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_13, _mma_ss_b_desc_13, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_13, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_13, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_13, _mma_ss_b_desc_13, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_14 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_14 = make_warp_uniform((((smem_q3_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_14 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_14); + uint64_t _mma_ss_b_desc_14 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_14); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_14, _mma_ss_b_desc_14, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_14, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_14, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_14, _mma_ss_b_desc_14, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_14, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_14, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_14, _mma_ss_b_desc_14, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_14, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_14, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_14, _mma_ss_b_desc_14, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_15 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_15 = make_warp_uniform((((smem_q4_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_15 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_15); + uint64_t _mma_ss_b_desc_15 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_15); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_15, _mma_ss_b_desc_15, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_15, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_15, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_15, _mma_ss_b_desc_15, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_15, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_15, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_15, _mma_ss_b_desc_15, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_15, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_15, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_15, _mma_ss_b_desc_15, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + elect_commit(s_full0_addr); + mbarrier_wait(s_empty1_addr, _phase_s_empty1_0); + _phase_s_empty1_0 ^= 1; + mbarrier_wait(o_empty1_addr, _phase_o_empty1_0); + _phase_o_empty1_0 ^= 1; + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_16 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (kv_stage_m) * 1024); + int _mma_b_lo_16 = make_warp_uniform((((smem_p1_addr) >> 4) & 0x3FFF) | 0x800000); + { + uint64_t _mma_ss_a_desc_16 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_16); + uint64_t _mma_ss_b_desc_16 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_16); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_16, _mma_ss_b_desc_16, 134512656, ((pair_m == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_16, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_16, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_16, _mma_ss_b_desc_16, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_16, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_16, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_16, _mma_ss_b_desc_16, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_16, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_16, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_16, _mma_ss_b_desc_16, 134512656, 1); + } + } + elect_commit2(kv_empty_addr + (kv_stage_m) * 8, o_full1_addr); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_17 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_17 = make_warp_uniform((((smem_q0_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_17 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_17); + uint64_t _mma_ss_b_desc_17 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_17); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_17, _mma_ss_b_desc_17, 134479888, 0); + } + incr_smem_desc_lo(_mma_ss_a_desc_17, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_17, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_17, _mma_ss_b_desc_17, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_17, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_17, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_17, _mma_ss_b_desc_17, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_17, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_17, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_17, _mma_ss_b_desc_17, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_18 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_18 = make_warp_uniform((((smem_q1_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_18 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_18); + uint64_t _mma_ss_b_desc_18 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_18); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_18, _mma_ss_b_desc_18, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_18, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_18, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_18, _mma_ss_b_desc_18, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_18, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_18, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_18, _mma_ss_b_desc_18, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_18, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_18, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_18, _mma_ss_b_desc_18, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_19 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_19 = make_warp_uniform((((smem_q2_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_19 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_19); + uint64_t _mma_ss_b_desc_19 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_19); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_19, _mma_ss_b_desc_19, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_19, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_19, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_19, _mma_ss_b_desc_19, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_19, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_19, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_19, _mma_ss_b_desc_19, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_19, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_19, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_19, _mma_ss_b_desc_19, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_20 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_20 = make_warp_uniform((((smem_q3_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_20 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_20); + uint64_t _mma_ss_b_desc_20 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_20); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_20, _mma_ss_b_desc_20, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_20, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_20, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_20, _mma_ss_b_desc_20, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_20, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_20, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_20, _mma_ss_b_desc_20, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_20, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_20, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_20, _mma_ss_b_desc_20, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_21 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_21 = make_warp_uniform((((smem_q4_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_21 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_21); + uint64_t _mma_ss_b_desc_21 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_21); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_21, _mma_ss_b_desc_21, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_21, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_21, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_21, _mma_ss_b_desc_21, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_21, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_21, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_21, _mma_ss_b_desc_21, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_21, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_21, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_21, _mma_ss_b_desc_21, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + elect_commit(s_full1_addr); + } + mbarrier_wait(s_empty0_addr, _phase_s_empty0_0); + _phase_s_empty0_0 ^= 1; + mbarrier_wait(o_empty0_addr, _phase_o_empty0_0); + _phase_o_empty0_0 ^= 1; + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_22 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (kv_stage_m) * 1024); + int _mma_b_lo_22 = make_warp_uniform((((smem_p0_addr) >> 4) & 0x3FFF) | 0x800000); + { + uint64_t _mma_ss_a_desc_22 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_22); + uint64_t _mma_ss_b_desc_22 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_22); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_22, _mma_ss_b_desc_22, 134512656, ((num_steps_m == 1) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_22, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_22, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_22, _mma_ss_b_desc_22, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_22, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_22, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_22, _mma_ss_b_desc_22, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_22, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_22, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_22, _mma_ss_b_desc_22, 134512656, 1); + } + } + elect_commit2(kv_empty_addr + (kv_stage_m) * 8, o_full0_addr); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(s_empty1_addr, _phase_s_empty1_0); + _phase_s_empty1_0 ^= 1; + mbarrier_wait(o_empty1_addr, _phase_o_empty1_0); + _phase_o_empty1_0 ^= 1; + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_23 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (kv_stage_m) * 1024); + int _mma_b_lo_23 = make_warp_uniform((((smem_p1_addr) >> 4) & 0x3FFF) | 0x800000); + { + uint64_t _mma_ss_a_desc_23 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_23); + uint64_t _mma_ss_b_desc_23 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_23); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_23, _mma_ss_b_desc_23, 134512656, ((num_steps_m == 1) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_23, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_23, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_23, _mma_ss_b_desc_23, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_23, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_23, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_23, _mma_ss_b_desc_23, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_23, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_23, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_23, _mma_ss_b_desc_23, 134512656, 1); + } + } + elect_commit2(kv_empty_addr + (kv_stage_m) * 8, o_full1_addr); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + elect_commit(q_empty_addr); + } + // ---- Role: page_warp ---- + } else if (warp == 13) { + { // page_warp_main + int cta_q_p = blockIdx.x / 2; + int cta_kv_p = blockIdx.x % 2; + int batch_p = blockIdx.z; + int seq_len_p = seq_lens[batch_p] - (q_len - 1 - cta_q_p); + int num_ctas_p = 2; + int num_steps_p = 2; + int num_tiles_p = num_steps_p * 2; + int max_page_p = (seq_len_p + 32 - 1) / 32 - 1; + int page_stage_p = 0; + unsigned int _phase_page_empty = 1; + #pragma unroll 1 + for (int tile_p = 0; tile_p < num_tiles_p; tile_p++) { + mbarrier_wait(page_empty_addr + (page_stage_p) * 8, _phase_page_empty); + int global_page_p = (unsigned int)((cta_kv_p * num_tiles_p + tile_p) * 4) + lane; + if (global_page_p > max_page_p) { + global_page_p = max_page_p; + } + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 4;" + :: "r"(page_offsets_addr + (unsigned int)(page_stage_p * 128) + lane * 4), "l"(page_table + (batch_p * page_table_stride + global_page_p))); + asm volatile( + "{\n\t" + "cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n\t" + "}" + :: "r"(page_full_addr + (page_stage_p) * 8) : "memory"); + page_stage_p += 1; + if (page_stage_p == 6) { page_stage_p = 0; _phase_page_empty ^= 1; } + } + } + // ---- Role: padding_warp ---- + } else if (warp == 14) { + // idle — no tasks assigned + // ---- Role: load_warp ---- + } else if (warp == 15) { + { // load_warp_main + int cta_q_l = blockIdx.x / 2; + int cta_kv_l = blockIdx.x % 2; + int y_l = blockIdx.y; + int head_group_l = y_l / 4; + int hv_partition_l = y_l % 4; + int batch_l = blockIdx.z; + int seq_len_l = seq_lens[batch_l] - (q_len - 1 - cta_q_l); + int num_ctas_l = 2; + int num_steps_l = 2; + int num_tiles_l = num_steps_l * 2; + int q_stage_l = 0; + int kv_stage_l = 0; + int page_stage_l = 0; + unsigned int _phase_q_empty = 1; + mbarrier_wait(q_empty_addr + (q_stage_l) * 8, _phase_q_empty); + if (elect_sync()) { + int q_row_l = (batch_l * q_len + cta_q_l) * 128 + head_group_l * 16; + mbarrier_arrive_expect_tx(q_full_addr + (q_stage_l) * 8, 10240); + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;" + :: "r"(smem_q0_addr + (unsigned int)(q_stage_l * 10240)), "l"((&Q_map)), "r"(0), "r"(q_row_l), + "r"(q_full_addr + (q_stage_l) * 8), "l"(0x14F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;" + :: "r"(smem_q1_addr + (unsigned int)(q_stage_l * 10240)), "l"((&Q_map)), "r"(128), "r"(q_row_l), + "r"(q_full_addr + (q_stage_l) * 8), "l"(0x14F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;" + :: "r"(smem_q2_addr + (unsigned int)(q_stage_l * 10240)), "l"((&Q_map)), "r"(256), "r"(q_row_l), + "r"(q_full_addr + (q_stage_l) * 8), "l"(0x14F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;" + :: "r"(smem_q3_addr + (unsigned int)(q_stage_l * 10240)), "l"((&Q_map)), "r"(384), "r"(q_row_l), + "r"(q_full_addr + (q_stage_l) * 8), "l"(0x14F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;" + :: "r"(smem_q4_addr + (unsigned int)(q_stage_l * 10240)), "l"((&Q_map)), "r"(512), "r"(q_row_l), + "r"(q_full_addr + (q_stage_l) * 8), "l"(0x14F0000000000000ULL) : "memory"); + } + unsigned int _phase_kv_empty = 1; + #pragma unroll 1 + for (int aligned_tile_l = 0; aligned_tile_l < num_tiles_l; aligned_tile_l++) { + mbarrier_wait(page_full_addr + (aligned_tile_l) * 8, 0); + #pragma unroll + for (int head_stage = 0; head_stage < 5; head_stage++) { + mbarrier_wait(kv_empty_addr + (kv_stage_l) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_l) * 8, 16384); + #pragma unroll + for (int page = 0; page < 4; page++) { + int physical_page = page_offsets[aligned_tile_l * 32 + page]; + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;" + :: "r"(smem_kv_addr + (unsigned int)(kv_stage_l * 16384) + (unsigned int)(page * 32 * 128)), "l"((&K_map)), "r"(head_stage * 128), "r"(physical_page * 32), + "r"(kv_full_addr + (kv_stage_l) * 8), "l"(0x14F0000000000000ULL) : "memory"); + } + } + kv_stage_l += 1; + if (kv_stage_l == 9) { kv_stage_l = 0; _phase_kv_empty ^= 1; } + } + if (aligned_tile_l > 0) { + mbarrier_wait(kv_empty_addr + (kv_stage_l) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_l) * 8, 16384); + #pragma unroll + for (int page_1 = 0; page_1 < 4; page_1++) { + int physical_page_1 = page_offsets[(aligned_tile_l - 1) * 32 + page_1]; + tma_2d_gmem2smem(smem_kv_addr + (unsigned int)(kv_stage_l * 16384) + (unsigned int)(page_1 * 32 * 128), (&K_map), hv_partition_l * 128, physical_page_1 * 32, kv_full_addr + (kv_stage_l) * 8); + } + } + kv_stage_l += 1; + if (kv_stage_l == 9) { kv_stage_l = 0; _phase_kv_empty ^= 1; } + mbarrier_arrive(page_empty_addr + (aligned_tile_l - 1) * 8); + } + } + mbarrier_wait(kv_empty_addr + (kv_stage_l) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_l) * 8, 16384); + #pragma unroll + for (int page_2 = 0; page_2 < 4; page_2++) { + int physical_page_2 = page_offsets[(num_tiles_l - 1) * 32 + page_2]; + tma_2d_gmem2smem(smem_kv_addr + (unsigned int)(kv_stage_l * 16384) + (unsigned int)(page_2 * 32 * 128), (&K_map), hv_partition_l * 128, physical_page_2 * 32, kv_full_addr + (kv_stage_l) * 8); + } + } + kv_stage_l += 1; + if (kv_stage_l == 9) { kv_stage_l = 0; _phase_kv_empty ^= 1; } + mbarrier_arrive(page_empty_addr + (num_tiles_l - 1) * 8); + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_full_abi_tail_bf16_83451b7317.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_full_abi_tail_bf16_83451b7317.cu new file mode 100644 index 00000000000..c2a6c26e87a --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_full_abi_tail_bf16_83451b7317.cu @@ -0,0 +1,196 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_WEIGHTS_OFF 0 +#define SMEM_WEIGHTS_STAGE_BYTES 16384 +#define SMEM_WEIGHTS_STRIDE 16384 +#define SMEM_TOTAL 16384 +#define THREADS 32 + +#include + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + +extern "C" { + +__global__ __launch_bounds__(32, 1) void +kernel_full_abi_tail_bf16(__nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ KV, float* __restrict__ fp8_lut, int* __restrict__ page_table, int* __restrict__ sparse_indices, int* __restrict__ row_batches, int* __restrict__ row_seq_lens, __nv_bfloat16* __restrict__ O, float* __restrict__ LSE, float* __restrict__ sinks, int num_heads, int qk_dim, int value_dim, int kv_stride, int page_size, int page_table_width, int sparse_width, int use_sparse, float softmax_scale, float bmm2_scale, int enable_sink, int write_lse) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + float* weights = reinterpret_cast(smem_raw + 0); + const int weights_addr = smem + 0; + + // === Task calls (dependency order) === + int row = blockIdx.x; + int head = blockIdx.y; + int batch = row_batches[row]; + int seq_len = row_seq_lens[row]; + int q_base = (row * num_heads + head) * qk_dim; + float local_max = -MLA_INF; + #pragma unroll 1 + for (int token = lane; token < seq_len; token += 32) { + int physical_token = 0; + if (use_sparse != 0) { + physical_token = sparse_indices[row * sparse_width + token]; + } else { + int logical_page = token / page_size; + int token_in_page = token - logical_page * page_size; + int physical_page = page_table[batch * page_table_width + logical_page]; + physical_token = physical_page * page_size + token_in_page; + } + int kv_base = physical_token * kv_stride; + float score = 0.0f; + #pragma unroll 4 + for (int dim = 0; dim < qk_dim; dim++) { + float q_value = Q[q_base + dim]; + float kv_value = KV[kv_base + dim]; + score = score + q_value * kv_value; + } + score = score * softmax_scale; + weights[token] = score; + float _max_0 = max_noftz(local_max, score); + local_max = _max_0; + } + if (enable_sink != 0 && lane == 0) { + float _max_1 = max_noftz(local_max, sinks[head]); + local_max = _max_1; + } + float _warp_reduce_0 = local_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + float row_max = _warp_reduce_0; + __syncthreads(); + float local_sum = 0.0f; + #pragma unroll 1 + for (int token_1 = lane; token_1 < seq_len; token_1 += 32) { + float _exp2_0 = approx_exp2((weights[token_1] - row_max) * 1.4426950408889634f); + float weight = _exp2_0; + weights[token_1] = weight; + local_sum = local_sum + weight; + } + if (enable_sink != 0 && lane == 0) { + float _exp2_1 = approx_exp2((sinks[head] - row_max) * 1.4426950408889634f); + local_sum = local_sum + _exp2_1; + } + float _warp_reduce_1 = local_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + float row_sum = _warp_reduce_1; + float _rcp_0 = approx_rcp(row_sum); + float inv_sum = _rcp_0; + __syncthreads(); + int o_base = (row * num_heads + head) * value_dim; + #pragma unroll 1 + for (int dim_1 = lane; dim_1 < value_dim; dim_1 += 32) { + float accum = 0.0f; + #pragma unroll 1 + for (int token_2 = 0; token_2 < seq_len; token_2++) { + int physical_token_1 = 0; + if (use_sparse != 0) { + physical_token_1 = sparse_indices[row * sparse_width + token_2]; + } else { + int logical_page_1 = token_2 / page_size; + int token_in_page_1 = token_2 - logical_page_1 * page_size; + int physical_page_1 = page_table[batch * page_table_width + logical_page_1]; + physical_token_1 = physical_page_1 * page_size + token_in_page_1; + } + int kv_base_1 = physical_token_1 * kv_stride; + float value = KV[kv_base_1 + dim_1]; + accum = accum + weights[token_2] * value; + } + *(reinterpret_cast<__nv_bfloat16*>(O + (o_base + dim_1)) + (0)) = __float2bfloat16_rn(accum * inv_sum * bmm2_scale); + } + if (write_lse != 0 && lane == 0) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(row_sum)); + float lse = row_max * 1.4426950408889634f + _log2_0; + *(reinterpret_cast(LSE + (row * num_heads + head)) + (0)) = lse; + } +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_full_abi_tail_fp8_26b36022eb.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_full_abi_tail_fp8_26b36022eb.cu new file mode 100644 index 00000000000..4e914cabd81 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_full_abi_tail_fp8_26b36022eb.cu @@ -0,0 +1,196 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_WEIGHTS_OFF 0 +#define SMEM_WEIGHTS_STAGE_BYTES 16384 +#define SMEM_WEIGHTS_STRIDE 16384 +#define SMEM_TOTAL 16384 +#define THREADS 32 + +#include + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + +extern "C" { + +__global__ __launch_bounds__(32, 1) void +kernel_full_abi_tail_fp8(uint8_t* __restrict__ Q, uint8_t* __restrict__ KV, float* __restrict__ fp8_lut, int* __restrict__ page_table, int* __restrict__ sparse_indices, int* __restrict__ row_batches, int* __restrict__ row_seq_lens, __nv_bfloat16* __restrict__ O, float* __restrict__ LSE, float* __restrict__ sinks, int num_heads, int qk_dim, int value_dim, int kv_stride, int page_size, int page_table_width, int sparse_width, int use_sparse, float softmax_scale, float bmm2_scale, int enable_sink, int write_lse) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + float* weights = reinterpret_cast(smem_raw + 0); + const int weights_addr = smem + 0; + + // === Task calls (dependency order) === + int row = blockIdx.x; + int head = blockIdx.y; + int batch = row_batches[row]; + int seq_len = row_seq_lens[row]; + int q_base = (row * num_heads + head) * qk_dim; + float local_max = -MLA_INF; + #pragma unroll 1 + for (int token = lane; token < seq_len; token += 32) { + int physical_token = 0; + if (use_sparse != 0) { + physical_token = sparse_indices[row * sparse_width + token]; + } else { + int logical_page = token / page_size; + int token_in_page = token - logical_page * page_size; + int physical_page = page_table[batch * page_table_width + logical_page]; + physical_token = physical_page * page_size + token_in_page; + } + int kv_base = physical_token * kv_stride; + float score = 0.0f; + #pragma unroll 4 + for (int dim = 0; dim < qk_dim; dim++) { + float q_value = fp8_lut[Q[q_base + dim]]; + float kv_value = fp8_lut[KV[kv_base + dim]]; + score = score + q_value * kv_value; + } + score = score * softmax_scale; + weights[token] = score; + float _max_0 = max_noftz(local_max, score); + local_max = _max_0; + } + if (enable_sink != 0 && lane == 0) { + float _max_1 = max_noftz(local_max, sinks[head]); + local_max = _max_1; + } + float _warp_reduce_0 = local_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + float row_max = _warp_reduce_0; + __syncthreads(); + float local_sum = 0.0f; + #pragma unroll 1 + for (int token_1 = lane; token_1 < seq_len; token_1 += 32) { + float _exp2_0 = approx_exp2((weights[token_1] - row_max) * 1.4426950408889634f); + float weight = _exp2_0; + weights[token_1] = weight; + local_sum = local_sum + weight; + } + if (enable_sink != 0 && lane == 0) { + float _exp2_1 = approx_exp2((sinks[head] - row_max) * 1.4426950408889634f); + local_sum = local_sum + _exp2_1; + } + float _warp_reduce_1 = local_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + float row_sum = _warp_reduce_1; + float _rcp_0 = approx_rcp(row_sum); + float inv_sum = _rcp_0; + __syncthreads(); + int o_base = (row * num_heads + head) * value_dim; + #pragma unroll 1 + for (int dim_1 = lane; dim_1 < value_dim; dim_1 += 32) { + float accum = 0.0f; + #pragma unroll 1 + for (int token_2 = 0; token_2 < seq_len; token_2++) { + int physical_token_1 = 0; + if (use_sparse != 0) { + physical_token_1 = sparse_indices[row * sparse_width + token_2]; + } else { + int logical_page_1 = token_2 / page_size; + int token_in_page_1 = token_2 - logical_page_1 * page_size; + int physical_page_1 = page_table[batch * page_table_width + logical_page_1]; + physical_token_1 = physical_page_1 * page_size + token_in_page_1; + } + int kv_base_1 = physical_token_1 * kv_stride; + float value = fp8_lut[KV[kv_base_1 + dim_1]]; + accum = accum + weights[token_2] * value; + } + *(reinterpret_cast<__nv_bfloat16*>(O + (o_base + dim_1)) + (0)) = __float2bfloat16_rn(accum * inv_sum * bmm2_scale); + } + if (write_lse != 0 && lane == 0) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(row_sum)); + float lse = row_max * 1.4426950408889634f + _log2_0; + *(reinterpret_cast(LSE + (row * num_heads + head)) + (0)) = lse; + } +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_30cfe17e1f.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_30cfe17e1f.cu new file mode 100644 index 00000000000..a8869605e0a --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_30cfe17e1f.cu @@ -0,0 +1,2469 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ uint32_t mbarrier_test_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.test_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_test_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 80;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < total_work_items; work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (total_work_items > 74) { + if (total_work_items <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * max_pages_per_seq; + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)value_split_count + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < total_work_items; _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)value_split_count + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < total_work_items; work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (total_work_items > 74) { + if (total_work_items <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (int)(((unsigned int)seqlen_kv_b_1 + 128 - 1) / 128); + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + uint32_t _mbar_token_10 = mbarrier_test_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + int abs_tile = start_tile_1 + tile_1; + int tile_end_k = (abs_tile + 1) * 128; + bool score_tile_full = tile_end_k <= seqlen_kv_b_1; + float sv[64]; + float tile_max = -MLA_INF; + float new_max = -MLA_INF; + if (!score_tile_full) { + mbarrier_wait_token(s_full_addr + (phase) * 8, s_wait_phase, _mbar_token_10); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } else { + mbarrier_wait_token(s_full_addr + (phase) * 8, s_wait_phase, _mbar_token_10); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + } + { + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + } + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)value_split_count + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 216;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < total_work_items; work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (total_work_items > 74) { + if (total_work_items <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + { + { + #pragma unroll 1 + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll 1 + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + const float2 _scale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_8)[_ls] = mul_f32x2_rn_noftz(reinterpret_cast(_tmem_load_8)[_ls], _scale2_1); + unsigned int pair_packed[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_8[_lp*2 + 0], _tmem_load_8[_lp*2+1 + 0])); + pair_packed[_lp] = *(uint32_t*)&_bf2; + } + int pair_parity = lane & 1; + unsigned int pair_mask = (unsigned int)(-pair_parity); + int pair_neighbor = lane ^ 1; + int pair_even_base = gmem_base - pair_parity * 512 + pair_parity * 8; + #pragma unroll + for (int pair_columns = 0; pair_columns < 2; pair_columns++) { + unsigned int pair_even_words[4]; + unsigned int pair_odd_words[4]; + #pragma unroll + for (int pair_word = 0; pair_word < 4; pair_word++) { + unsigned int _shfl_0 = __shfl_sync(0xFFFFFFFF, pair_packed[pair_columns * 8 + 4 + pair_word] & ~pair_mask | pair_packed[pair_columns * 8 + pair_word] & pair_mask, pair_neighbor); + pair_even_words[pair_word] = pair_packed[pair_columns * 8 + pair_word] & ~pair_mask | _shfl_0 & pair_mask; + pair_odd_words[pair_word] = _shfl_0 & ~pair_mask | pair_packed[pair_columns * 8 + 4 + pair_word] & pair_mask; + } + reinterpret_cast(O + (pair_even_base + pair_columns * 16))[0] = reinterpret_cast(pair_even_words)[0]; + reinterpret_cast(O + (pair_even_base + 512 + pair_columns * 16))[0] = reinterpret_cast(pair_odd_words)[0]; + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)value_split_count + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < total_work_items; work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (total_work_items > 74) { + if (total_work_items <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)value_split_count + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < total_work_items; _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)value_split_count + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_3b5bd1e042.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_3b5bd1e042.cu new file mode 100644 index 00000000000..4c815fa3ab0 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_3b5bd1e042.cu @@ -0,0 +1,2435 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 +#define Q8_B768 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < ((Q8_B768) ? 6144 : total_work_items); work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 74) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * ((Q8_B768) ? 32 : max_pages_per_seq); + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + if (n == 0) { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + if (n_1 == 0) { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < ((Q8_B768) ? 6144 : total_work_items); _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < ((Q8_B768) ? 6144 : total_work_items); work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 74) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < ((Q8_B768) ? 6144 : total_work_items); work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 74) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll 1 + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_8[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_8[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_8[j + 0], _tmem_load_8[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_8[j + 2], _tmem_load_8[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_8[j + 4], _tmem_load_8[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_8[j + 6], _tmem_load_8[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < ((Q8_B768) ? 6144 : total_work_items); work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 74) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < ((Q8_B768) ? 6144 : total_work_items); _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_7aa2c0c00a.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_7aa2c0c00a.cu new file mode 100644 index 00000000000..11e0f3c0ae8 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_7aa2c0c00a.cu @@ -0,0 +1,2490 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 1024 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 1024 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 230400 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230912 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 231424 +#define THREADS 384 +#define Q4_B1024 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x8(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x8.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem + 230944; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 230400); + const int smem_epilogue_exchange_addr = smem + 230400; + unsigned int* work_response = reinterpret_cast(smem_raw + 230912); + const int work_response_addr = smem + 230912; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230944..231392) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230944 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 231200 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 231392); + if (warp == 0) { + int _tmem_hold = smem + 231392; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * ((Q4_B1024) ? 4 : value_split_count) + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < ((Q4_B1024) ? 4096 : total_work_items); work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (((Q4_B1024) ? 4096 : total_work_items) > 74) { + if (((Q4_B1024) ? 4096 : total_work_items) <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[((Q4_B1024) ? 4096 : total_work_items) * ((Q4_B1024) ? 32 : max_pages_per_seq) + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * ((Q4_B1024) ? 32 : max_pages_per_seq); + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)(((Q4_B1024) ? 4 : value_split_count)) + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * ((Q4_B1024) ? 4 : value_split_count) + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < ((Q4_B1024) ? 4096 : total_work_items); _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)(((Q4_B1024) ? 4 : value_split_count)) + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 176;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * ((Q4_B1024) ? 4 : value_split_count) + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < ((Q4_B1024) ? 4096 : total_work_items); work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (((Q4_B1024) ? 4096 : total_work_items) > 74) { + if (((Q4_B1024) ? 4096 : total_work_items) <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[((Q4_B1024) ? 4096 : total_work_items) * ((Q4_B1024) ? 32 : max_pages_per_seq) + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[phase * 128 + exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[phase * 128 + exchange_idx ^ 64]); + new_max = _max_1; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)(((Q4_B1024) ? 4 : value_split_count)) + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * ((Q4_B1024) ? 4 : value_split_count) + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < ((Q4_B1024) ? 4096 : total_work_items); work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (((Q4_B1024) ? 4096 : total_work_items) > 74) { + if (((Q4_B1024) ? 4096 : total_work_items) <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[((Q4_B1024) ? 4096 : total_work_items) * ((Q4_B1024) ? 32 : max_pages_per_seq) + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll 1 + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll 1 + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll 1 + for (int c_1 = 0; c_1 < 64; c_1 += 64) { + float _tmem_load_6[32]; + tmem_ld_x8(&_tmem_load_6[0], o_base_epi + c_1); + tmem_ld_x8(&_tmem_load_6[8], o_base_epi + c_1 + 8); + tmem_ld_x8(&_tmem_load_6[16], o_base_epi + c_1 + 16); + tmem_ld_x8(&_tmem_load_6[24], o_base_epi + c_1 + 24); + float _tmem_load_7[32]; + tmem_ld_x8(&_tmem_load_7[0], o_base_epi + c_1 + 32); + tmem_ld_x8(&_tmem_load_7[8], o_base_epi + c_1 + 32 + 8); + tmem_ld_x8(&_tmem_load_7[16], o_base_epi + c_1 + 32 + 16); + tmem_ld_x8(&_tmem_load_7[24], o_base_epi + c_1 + 32 + 24); + unsigned int epilogue_packed_first[16]; + unsigned int epilogue_packed_next[16]; + const float2 _scale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_6)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_6)[_ls], _scale2_1); + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_6[_lp*2 + 0], _tmem_load_6[_lp*2+1 + 0])); + epilogue_packed_first[_lp] = *(uint32_t*)&_bf2; + } + const float2 _scale2_2 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_7)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_7)[_ls], _scale2_2); + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_7[_lp*2 + 0], _tmem_load_7[_lp*2+1 + 0])); + epilogue_packed_next[_lp] = *(uint32_t*)&_bf2; + } + int gmem_base = o_offset + logical_vs * 64 + c_1; + if (((unsigned long long)O & 31) == 0) { + #pragma unroll + for (int j = 0; j < 32; j += 16) { + { + unsigned _stv8_3_0 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 0]); + unsigned _stv8_3_1 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 1]); + unsigned _stv8_3_2 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 2]); + unsigned _stv8_3_3 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 3]); + unsigned _stv8_3_4 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 4]); + unsigned _stv8_3_5 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 5]); + unsigned _stv8_3_6 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 6]); + unsigned _stv8_3_7 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (gmem_base + j) + (0))), "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3), "r"(_stv8_3_4), "r"(_stv8_3_5), "r"(_stv8_3_6), "r"(_stv8_3_7) : "memory"); + } + } + #pragma unroll + for (int j_1 = 0; j_1 < 32; j_1 += 16) { + { + unsigned _stv8_4_0 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 0]); + unsigned _stv8_4_1 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 1]); + unsigned _stv8_4_2 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 2]); + unsigned _stv8_4_3 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 3]); + unsigned _stv8_4_4 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 4]); + unsigned _stv8_4_5 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 5]); + unsigned _stv8_4_6 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 6]); + unsigned _stv8_4_7 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (gmem_base + 32 + j_1) + (0))), "r"(_stv8_4_0), "r"(_stv8_4_1), "r"(_stv8_4_2), "r"(_stv8_4_3), "r"(_stv8_4_4), "r"(_stv8_4_5), "r"(_stv8_4_6), "r"(_stv8_4_7) : "memory"); + } + } + } else { + #pragma unroll + for (int j_2 = 0; j_2 < 32; j_2 += 8) { + reinterpret_cast(O + (gmem_base + j_2))[0] = reinterpret_cast(epilogue_packed_first + j_2 / 2)[0]; + } + #pragma unroll + for (int j_3 = 0; j_3 < 32; j_3 += 8) { + reinterpret_cast(O + (gmem_base + 32 + j_3))[0] = reinterpret_cast(epilogue_packed_next + j_3 / 2)[0]; + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)(((Q4_B1024) ? 4 : value_split_count)) + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * ((Q4_B1024) ? 4 : value_split_count) + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < ((Q4_B1024) ? 4096 : total_work_items); work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (((Q4_B1024) ? 4096 : total_work_items) > 74) { + if (((Q4_B1024) ? 4096 : total_work_items) <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[((Q4_B1024) ? 4096 : total_work_items) * ((Q4_B1024) ? 32 : max_pages_per_seq) + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)(((Q4_B1024) ? 4 : value_split_count)) + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * ((Q4_B1024) ? 4 : value_split_count) + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < ((Q4_B1024) ? 4096 : total_work_items); _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)(((Q4_B1024) ? 4 : value_split_count)) + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_8587a85908.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_8587a85908.cu new file mode 100644 index 00000000000..a8dcd54a2a5 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_8587a85908.cu @@ -0,0 +1,2490 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 1024 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 1024 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 230400 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230912 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 231424 +#define THREADS 384 +#define EXACT007826_CONSTANTS 0 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x8(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x8.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem + 230944; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 230400); + const int smem_epilogue_exchange_addr = smem + 230400; + unsigned int* work_response = reinterpret_cast(smem_raw + 230912); + const int work_response_addr = smem + 230912; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230944..231392) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230944 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 231200 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 231392); + if (warp == 0) { + int _tmem_hold = smem + 231392; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * ((EXACT007826_CONSTANTS) ? 16 : value_split_count) + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < ((EXACT007826_CONSTANTS) ? 8192 : total_work_items); work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) > 74) { + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[((EXACT007826_CONSTANTS) ? 8192 : total_work_items) * ((EXACT007826_CONSTANTS) ? 32 : max_pages_per_seq) + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * ((EXACT007826_CONSTANTS) ? 32 : max_pages_per_seq); + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)(((EXACT007826_CONSTANTS) ? 16 : value_split_count)) + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * ((EXACT007826_CONSTANTS) ? 16 : value_split_count) + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < ((EXACT007826_CONSTANTS) ? 8192 : total_work_items); _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)(((EXACT007826_CONSTANTS) ? 16 : value_split_count)) + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 184;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * ((EXACT007826_CONSTANTS) ? 16 : value_split_count) + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < ((EXACT007826_CONSTANTS) ? 8192 : total_work_items); work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) > 74) { + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[((EXACT007826_CONSTANTS) ? 8192 : total_work_items) * ((EXACT007826_CONSTANTS) ? 32 : max_pages_per_seq) + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[phase * 128 + exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[phase * 128 + exchange_idx ^ 64]); + new_max = _max_1; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)(((EXACT007826_CONSTANTS) ? 16 : value_split_count)) + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * ((EXACT007826_CONSTANTS) ? 16 : value_split_count) + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < ((EXACT007826_CONSTANTS) ? 8192 : total_work_items); work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) > 74) { + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[((EXACT007826_CONSTANTS) ? 8192 : total_work_items) * ((EXACT007826_CONSTANTS) ? 32 : max_pages_per_seq) + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll 1 + for (int c_1 = 0; c_1 < 64; c_1 += 64) { + float _tmem_load_6[32]; + tmem_ld_x8(&_tmem_load_6[0], o_base_epi + c_1); + tmem_ld_x8(&_tmem_load_6[8], o_base_epi + c_1 + 8); + tmem_ld_x8(&_tmem_load_6[16], o_base_epi + c_1 + 16); + tmem_ld_x8(&_tmem_load_6[24], o_base_epi + c_1 + 24); + float _tmem_load_7[32]; + tmem_ld_x8(&_tmem_load_7[0], o_base_epi + c_1 + 32); + tmem_ld_x8(&_tmem_load_7[8], o_base_epi + c_1 + 32 + 8); + tmem_ld_x8(&_tmem_load_7[16], o_base_epi + c_1 + 32 + 16); + tmem_ld_x8(&_tmem_load_7[24], o_base_epi + c_1 + 32 + 24); + unsigned int epilogue_packed_first[16]; + unsigned int epilogue_packed_next[16]; + const float2 _scale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_6)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_6)[_ls], _scale2_1); + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_6[_lp*2 + 0], _tmem_load_6[_lp*2+1 + 0])); + epilogue_packed_first[_lp] = *(uint32_t*)&_bf2; + } + const float2 _scale2_2 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_7)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_7)[_ls], _scale2_2); + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_7[_lp*2 + 0], _tmem_load_7[_lp*2+1 + 0])); + epilogue_packed_next[_lp] = *(uint32_t*)&_bf2; + } + int gmem_base = o_offset + logical_vs * 64 + c_1; + if (((unsigned long long)O & 31) == 0) { + #pragma unroll + for (int j = 0; j < 32; j += 16) { + { + unsigned _stv8_3_0 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 0]); + unsigned _stv8_3_1 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 1]); + unsigned _stv8_3_2 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 2]); + unsigned _stv8_3_3 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 3]); + unsigned _stv8_3_4 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 4]); + unsigned _stv8_3_5 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 5]); + unsigned _stv8_3_6 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 6]); + unsigned _stv8_3_7 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (gmem_base + j) + (0))), "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3), "r"(_stv8_3_4), "r"(_stv8_3_5), "r"(_stv8_3_6), "r"(_stv8_3_7) : "memory"); + } + } + #pragma unroll + for (int j_1 = 0; j_1 < 32; j_1 += 16) { + { + unsigned _stv8_4_0 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 0]); + unsigned _stv8_4_1 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 1]); + unsigned _stv8_4_2 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 2]); + unsigned _stv8_4_3 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 3]); + unsigned _stv8_4_4 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 4]); + unsigned _stv8_4_5 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 5]); + unsigned _stv8_4_6 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 6]); + unsigned _stv8_4_7 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (gmem_base + 32 + j_1) + (0))), "r"(_stv8_4_0), "r"(_stv8_4_1), "r"(_stv8_4_2), "r"(_stv8_4_3), "r"(_stv8_4_4), "r"(_stv8_4_5), "r"(_stv8_4_6), "r"(_stv8_4_7) : "memory"); + } + } + } else { + #pragma unroll + for (int j_2 = 0; j_2 < 32; j_2 += 8) { + reinterpret_cast(O + (gmem_base + j_2))[0] = reinterpret_cast(epilogue_packed_first + j_2 / 2)[0]; + } + #pragma unroll + for (int j_3 = 0; j_3 < 32; j_3 += 8) { + reinterpret_cast(O + (gmem_base + 32 + j_3))[0] = reinterpret_cast(epilogue_packed_next + j_3 / 2)[0]; + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)(((EXACT007826_CONSTANTS) ? 16 : value_split_count)) + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * ((EXACT007826_CONSTANTS) ? 16 : value_split_count) + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < ((EXACT007826_CONSTANTS) ? 8192 : total_work_items); work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) > 74) { + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[((EXACT007826_CONSTANTS) ? 8192 : total_work_items) * ((EXACT007826_CONSTANTS) ? 32 : max_pages_per_seq) + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)(((EXACT007826_CONSTANTS) ? 16 : value_split_count)) + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * ((EXACT007826_CONSTANTS) ? 16 : value_split_count) + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < ((EXACT007826_CONSTANTS) ? 8192 : total_work_items); _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)(((EXACT007826_CONSTANTS) ? 16 : value_split_count)) + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_8ace9da7a5.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_8ace9da7a5.cu new file mode 100644 index 00000000000..347af73c6d5 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_8ace9da7a5.cu @@ -0,0 +1,2513 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 +#define Q8_B768 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < ((Q8_B768) ? 6144 : total_work_items); work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 74) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * ((Q8_B768) ? 32 : max_pages_per_seq); + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < ((Q8_B768) ? 6144 : total_work_items); _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < ((Q8_B768) ? 6144 : total_work_items); work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 74) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + float true_row_max_log2 = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + { + float candidate_max_log2 = new_max * softmax_scale_log2; + float normalization_max_log2 = row_max_val * softmax_scale_log2; + float _max_2 = max_noftz(true_row_max_log2, candidate_max_log2); + true_row_max_log2 = _max_2; + int _vote_0 = __any_sync(0xFFFFFFFF, candidate_max_log2 - normalization_max_log2 > 6.0f); + int rescale_vote = _vote_0; + if (rescale_vote == 0) { + new_max = row_max_val; + } + } + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < ((Q8_B768) ? 6144 : total_work_items); work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 74) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_1 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_1; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base); + float _tmem_load_3[16]; + tmem_ld_x16(&_tmem_load_3[0], o_base + 16); + float _tmem_load_4[16]; + tmem_ld_x16(&_tmem_load_4[0], o_base + 32); + float _tmem_load_5[16]; + tmem_ld_x16(&_tmem_load_5[0], o_base + 48); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base, _tmem_load_2); + const float2 _scale2_1 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_3)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_3)[_ls], _scale2_1); + tmem_st_x16_f32(o_base + 16, _tmem_load_3); + const float2 _scale2_2 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_4)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_4)[_ls], _scale2_2); + tmem_st_x16_f32(o_base + 32, _tmem_load_4); + const float2 _scale2_3 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_5)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_5)[_ls], _scale2_3); + tmem_st_x16_f32(o_base + 48, _tmem_load_5); + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_7[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_7[0]), "=f"(_tmem_load_7[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_7[0]; + final_max_val = _tmem_load_7[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c = 0; c < 64; c += 64) { + float _tmem_load_9[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_9[0]), "=f"(_tmem_load_9[1]), "=f"(_tmem_load_9[2]), "=f"(_tmem_load_9[3]), "=f"(_tmem_load_9[4]), "=f"(_tmem_load_9[5]), "=f"(_tmem_load_9[6]), "=f"(_tmem_load_9[7]), "=f"(_tmem_load_9[8]), "=f"(_tmem_load_9[9]), "=f"(_tmem_load_9[10]), "=f"(_tmem_load_9[11]), "=f"(_tmem_load_9[12]), "=f"(_tmem_load_9[13]), "=f"(_tmem_load_9[14]), "=f"(_tmem_load_9[15]), "=f"(_tmem_load_9[16]), "=f"(_tmem_load_9[17]), "=f"(_tmem_load_9[18]), "=f"(_tmem_load_9[19]), "=f"(_tmem_load_9[20]), "=f"(_tmem_load_9[21]), "=f"(_tmem_load_9[22]), "=f"(_tmem_load_9[23]), "=f"(_tmem_load_9[24]), "=f"(_tmem_load_9[25]), "=f"(_tmem_load_9[26]), "=f"(_tmem_load_9[27]), "=f"(_tmem_load_9[28]), "=f"(_tmem_load_9[29]), "=f"(_tmem_load_9[30]), "=f"(_tmem_load_9[31]) + : "r"(o_base_epi + c)); + float _tmem_load_10[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_10[0]), "=f"(_tmem_load_10[1]), "=f"(_tmem_load_10[2]), "=f"(_tmem_load_10[3]), "=f"(_tmem_load_10[4]), "=f"(_tmem_load_10[5]), "=f"(_tmem_load_10[6]), "=f"(_tmem_load_10[7]), "=f"(_tmem_load_10[8]), "=f"(_tmem_load_10[9]), "=f"(_tmem_load_10[10]), "=f"(_tmem_load_10[11]), "=f"(_tmem_load_10[12]), "=f"(_tmem_load_10[13]), "=f"(_tmem_load_10[14]), "=f"(_tmem_load_10[15]), "=f"(_tmem_load_10[16]), "=f"(_tmem_load_10[17]), "=f"(_tmem_load_10[18]), "=f"(_tmem_load_10[19]), "=f"(_tmem_load_10[20]), "=f"(_tmem_load_10[21]), "=f"(_tmem_load_10[22]), "=f"(_tmem_load_10[23]), "=f"(_tmem_load_10[24]), "=f"(_tmem_load_10[25]), "=f"(_tmem_load_10[26]), "=f"(_tmem_load_10[27]), "=f"(_tmem_load_10[28]), "=f"(_tmem_load_10[29]), "=f"(_tmem_load_10[30]), "=f"(_tmem_load_10[31]) + : "r"(o_base_epi + c + 32)); + unsigned int epilogue_packed_first[16]; + unsigned int epilogue_packed_next[16]; + const float2 _scale2_4 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_9)[_ls] = mul_f32x2_rn_noftz(reinterpret_cast(_tmem_load_9)[_ls], _scale2_4); + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_9[_lp*2 + 0], _tmem_load_9[_lp*2+1 + 0])); + epilogue_packed_first[_lp] = *(uint32_t*)&_bf2; + } + const float2 _scale2_5 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_10)[_ls] = mul_f32x2_rn_noftz(reinterpret_cast(_tmem_load_10)[_ls], _scale2_5); + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_10[_lp*2 + 0], _tmem_load_10[_lp*2+1 + 0])); + epilogue_packed_next[_lp] = *(uint32_t*)&_bf2; + } + int gmem_base = o_offset + logical_vs * 64 + c; + if (((unsigned long long)O & 31) == 0) { + #pragma unroll + for (int j = 0; j < 32; j += 16) { + { + unsigned _stv8_6_0 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 0]); + unsigned _stv8_6_1 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 1]); + unsigned _stv8_6_2 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 2]); + unsigned _stv8_6_3 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 3]); + unsigned _stv8_6_4 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 4]); + unsigned _stv8_6_5 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 5]); + unsigned _stv8_6_6 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 6]); + unsigned _stv8_6_7 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (gmem_base + j) + (0))), "r"(_stv8_6_0), "r"(_stv8_6_1), "r"(_stv8_6_2), "r"(_stv8_6_3), "r"(_stv8_6_4), "r"(_stv8_6_5), "r"(_stv8_6_6), "r"(_stv8_6_7) : "memory"); + } + } + #pragma unroll + for (int j_1 = 0; j_1 < 32; j_1 += 16) { + { + unsigned _stv8_7_0 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 0]); + unsigned _stv8_7_1 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 1]); + unsigned _stv8_7_2 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 2]); + unsigned _stv8_7_3 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 3]); + unsigned _stv8_7_4 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 4]); + unsigned _stv8_7_5 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 5]); + unsigned _stv8_7_6 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 6]); + unsigned _stv8_7_7 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (gmem_base + 32 + j_1) + (0))), "r"(_stv8_7_0), "r"(_stv8_7_1), "r"(_stv8_7_2), "r"(_stv8_7_3), "r"(_stv8_7_4), "r"(_stv8_7_5), "r"(_stv8_7_6), "r"(_stv8_7_7) : "memory"); + } + } + } else { + #pragma unroll + for (int j_2 = 0; j_2 < 32; j_2 += 8) { + reinterpret_cast(O + (gmem_base + j_2))[0] = reinterpret_cast(epilogue_packed_first + j_2 / 2)[0]; + } + #pragma unroll + for (int j_3 = 0; j_3 < 32; j_3 += 8) { + reinterpret_cast(O + (gmem_base + 32 + j_3))[0] = reinterpret_cast(epilogue_packed_next + j_3 / 2)[0]; + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < ((Q8_B768) ? 6144 : total_work_items); work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 74) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < ((Q8_B768) ? 6144 : total_work_items); _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_bbfb74c484.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_bbfb74c484.cu new file mode 100644 index 00000000000..8e21f908f28 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_bbfb74c484.cu @@ -0,0 +1,2435 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 +#define Q1_B512 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 80;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * ((Q1_B512) ? 1 : value_split_count) + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < ((Q1_B512) ? 512 : total_work_items); work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (((Q1_B512) ? 512 : total_work_items) > 74) { + if (((Q1_B512) ? 512 : total_work_items) <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[((Q1_B512) ? 512 : total_work_items) * max_pages_per_seq + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * max_pages_per_seq; + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)(((Q1_B512) ? 1 : value_split_count)) + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * ((Q1_B512) ? 1 : value_split_count) + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < ((Q1_B512) ? 512 : total_work_items); _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)(((Q1_B512) ? 1 : value_split_count)) + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * ((Q1_B512) ? 1 : value_split_count) + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < ((Q1_B512) ? 512 : total_work_items); work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (((Q1_B512) ? 512 : total_work_items) > 74) { + if (((Q1_B512) ? 512 : total_work_items) <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[((Q1_B512) ? 512 : total_work_items) * max_pages_per_seq + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)(((Q1_B512) ? 1 : value_split_count)) + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 216;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * ((Q1_B512) ? 1 : value_split_count) + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < ((Q1_B512) ? 512 : total_work_items); work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (((Q1_B512) ? 512 : total_work_items) > 74) { + if (((Q1_B512) ? 512 : total_work_items) <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[((Q1_B512) ? 512 : total_work_items) * max_pages_per_seq + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_8[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_8[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_8[j + 0], _tmem_load_8[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_8[j + 2], _tmem_load_8[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_8[j + 4], _tmem_load_8[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_8[j + 6], _tmem_load_8[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)(((Q1_B512) ? 1 : value_split_count)) + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * ((Q1_B512) ? 1 : value_split_count) + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < ((Q1_B512) ? 512 : total_work_items); work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (((Q1_B512) ? 512 : total_work_items) > 74) { + if (((Q1_B512) ? 512 : total_work_items) <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[((Q1_B512) ? 512 : total_work_items) * max_pages_per_seq + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)(((Q1_B512) ? 1 : value_split_count)) + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * ((Q1_B512) ? 1 : value_split_count) + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < ((Q1_B512) ? 512 : total_work_items); _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)(((Q1_B512) ? 1 : value_split_count)) + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_c278f217ca.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_c278f217ca.cu new file mode 100644 index 00000000000..916618946f4 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_c278f217ca.cu @@ -0,0 +1,2434 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < total_work_items; work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (total_work_items > 74) { + if (total_work_items <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * max_pages_per_seq; + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)value_split_count + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < total_work_items; _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)value_split_count + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < total_work_items; work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (total_work_items > 74) { + if (total_work_items <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)value_split_count + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < total_work_items; work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (total_work_items > 74) { + if (total_work_items <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_8[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_8[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_8[j + 0], _tmem_load_8[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_8[j + 2], _tmem_load_8[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_8[j + 4], _tmem_load_8[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_8[j + 6], _tmem_load_8[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)value_split_count + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < total_work_items; work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (total_work_items > 74) { + if (total_work_items <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)value_split_count + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < total_work_items; _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)value_split_count + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_c6609aef39.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_c6609aef39.cu new file mode 100644 index 00000000000..2e2aac1b84f --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_exact_live_bf16_clc_c6609aef39.cu @@ -0,0 +1,2434 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 88;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < total_work_items; work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (total_work_items > 74) { + if (total_work_items <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * max_pages_per_seq; + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)value_split_count + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < total_work_items; _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)value_split_count + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < total_work_items; work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (total_work_items > 74) { + if (total_work_items <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)value_split_count + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < total_work_items; work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (total_work_items > 74) { + if (total_work_items <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_8[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_8[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_8[j + 0], _tmem_load_8[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_8[j + 2], _tmem_load_8[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_8[j + 4], _tmem_load_8[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_8[j + 6], _tmem_load_8[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)value_split_count + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < total_work_items; work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (total_work_items > 74) { + if (total_work_items <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)value_split_count + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < total_work_items; _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)value_split_count + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_v32_paged_bf16_c80238a3c3.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_v32_paged_bf16_c80238a3c3.cu new file mode 100644 index 00000000000..b6995245cb8 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_v32_paged_bf16_c80238a3c3.cu @@ -0,0 +1,2073 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 74752 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_V_OFF 74752 +#define SMEM_SMEM_V_STAGE_BYTES 8192 +#define SMEM_SMEM_V_STRIDE 8192 +#define SMEM_SMEM_P_OFF 197632 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 230400 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 230912 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_TOTAL 231424 +#define THREADS 384 +#define USE_PDL 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384) __cluster_dims__(2,1,1) void +kernel_mla_decode_v32_paged_bf16(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k_nope, const __grid_constant__ CUtensorMap tmap_k_rope, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int num_split, int total_work_items, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_kv_addr = smem + 74752; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_v_addr = smem + 74752; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 197632); + const int smem_p_addr = smem + 197632; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 230400); + const int smem_softmax_exchange_addr = smem + 230400; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 230912); + const int smem_epilogue_exchange_addr = smem + 230912; + if (warp == 8) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); } + if (warp == 8) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k_nope))) : "memory"); } + if (warp == 8) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k_rope))) : "memory"); } + if (warp == 8) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); } + + // Mbarrier init (14 pipeline groups, 0 ordered-sequence groups, 48 barriers) + // Mbarriers at smem_raw[0..384) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + mbarrier_init(smem + 112, 1); + mbarrier_init(smem + 120, 1); + mbarrier_init(smem + 128, 1); + // kv_empty: 15 barriers, init_count=1 + mbarrier_init(smem + 136, 1); + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + mbarrier_init(smem + 160, 1); + mbarrier_init(smem + 168, 1); + mbarrier_init(smem + 176, 1); + mbarrier_init(smem + 184, 1); + mbarrier_init(smem + 192, 1); + mbarrier_init(smem + 200, 1); + mbarrier_init(smem + 208, 1); + mbarrier_init(smem + 216, 1); + mbarrier_init(smem + 224, 1); + mbarrier_init(smem + 232, 1); + mbarrier_init(smem + 240, 1); + mbarrier_init(smem + 248, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 256, 1); + mbarrier_init(smem + 264, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 272, 256); + mbarrier_init(smem + 280, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 288, 256); + mbarrier_init(smem + 296, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 304, 1); + mbarrier_init(smem + 312, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 320, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 328, 128); + mbarrier_init(smem + 336, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 344, 128); + mbarrier_init(smem + 352, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 360, 1); + // tmem_scrubbed: 1 barriers, init_count=256 + mbarrier_init(smem + 368, 256); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 376, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 384); + if (warp == 0) { + int _tmem_hold = smem + 384; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + #pragma unroll 1 + for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { + int batch_idx = work_idx / (unsigned int)num_split; + int split_idx = work_idx % (unsigned int)num_split; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int tiles_per_split = (total_kv_tiles + num_split - 1) / num_split; + int start_tile = split_idx * tiles_per_split; + int end_tile_raw = start_tile + tiles_per_split; + int end_tile = ((end_tile_raw < total_kv_tiles) ? end_tile_raw : total_kv_tiles); + int num_kv_tiles = end_tile - start_tile; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + #pragma unroll 1 + for (int tile = 0; tile < num_kv_tiles; tile++) { + int pipeline_tile = softmax_tile_cursor + tile; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(n_half * 64) + (unsigned int)(tmem_row_base << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile + tile; + int tail_valid = seqlen_kv_b - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -MLA_INF; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float _max_0 = max_noftz(_tmem_load_0_max, row_max_val); + smem_softmax_exchange[exchange_idx] = _max_0; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(_max_0, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float no_correction = (((_max_1 - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - _max_1); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = _max_1; + float safe_max = ((_max_1 == -MLA_INF) ? 0.0f : _max_1); + float max_scaled = safe_max * softmax_scale_log2; + #pragma unroll + for (int i = 0; i < 64; i++) { + float _exp2_1 = approx_exp2(_tmem_load_0[i] * softmax_scale_log2 - max_scaled); + _tmem_load_0[i] = _exp2_1; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_3); + softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_3); + float _tmem_load_0_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + _tmem_load_0_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (phase) * 8, p_cor_empty_phase); + float meta[4]; + meta[0] = row_sum_val; + meta[1] = row_max_val; + meta[2] = acc_scale; + meta[3] = no_correction; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + tmem_st_x4_f32(meta_addr, meta); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t _tmem_load_0_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(_tmem_load_0_bf16[pv_off])), "r"(__as_u32(_tmem_load_0_bf16[pv_off + 1])), "r"(__as_u32(_tmem_load_0_bf16[pv_off + 2])), "r"(__as_u32(_tmem_load_0_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { + int batch_idx_1 = work_idx_1 / (unsigned int)num_split; + int split_idx_1 = work_idx_1 % (unsigned int)num_split; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int tiles_per_split_1 = (total_kv_tiles_1 + num_split - 1) / num_split; + int start_tile_1 = split_idx_1 * tiles_per_split_1; + int end_tile_raw_1 = start_tile_1 + tiles_per_split_1; + int end_tile_1 = ((end_tile_raw_1 < total_kv_tiles_1) ? end_tile_raw_1 : total_kv_tiles_1); + int num_kv_tiles_1 = end_tile_1 - start_tile_1; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile_1 = correction_tile_cursor + tile_1; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], meta_addr_1); + final_sum_val = _tmem_load_1[0]; + final_max_val = _tmem_load_1[1]; + float acc_scale_1 = _tmem_load_1[2]; + float no_correction_1 = _tmem_load_1[3]; + if (tile_1 > 0) { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + if (tile_1 > 0) { + int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); + int skip_correction = _vote_0; + if (skip_correction == 0) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + #pragma unroll + for (int acc_stage = 0; acc_stage < 2; acc_stage++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + int stat_offset = (batch_idx_1 * 128 + head_idx) * num_split + split_idx_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + *(reinterpret_cast(partial_lse + stat_offset) + (0)) = lse_log2; + int o_offset = stat_offset * 512; + float output_scale = bmm2_scale; + if (enable_sink != 0 && num_split == 1) { + float _exp2_2 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - lse_log2); + float sink_exp = _exp2_2; + float _rcp_1 = approx_rcp(1.0f + sink_exp); + output_scale = output_scale * _rcp_1; + } + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_1 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_3[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_3[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_3[j + 0], _tmem_load_3[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_3[j + 2], _tmem_load_3[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_3[j + 4], _tmem_load_3[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_3[j + 6], _tmem_load_3[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_1; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + { + __threadfence(); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int wg2_dummy = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + unsigned int mma_kv_stage = 0; + int mma_tile_cursor = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + #pragma unroll 1 + for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { + int batch_idx_2 = work_idx_2 / (unsigned int)num_split; + int split_idx_2 = work_idx_2 % (unsigned int)num_split; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int tiles_per_split_2 = (total_kv_tiles_2 + num_split - 1) / num_split; + int start_tile_2 = split_idx_2 * tiles_per_split_2; + int end_tile_raw_2 = start_tile_2 + tiles_per_split_2; + int end_tile_2 = ((end_tile_raw_2 < total_kv_tiles_2) ? end_tile_raw_2 : total_kv_tiles_2); + int num_kv_tiles_2 = end_tile_2 - start_tile_2; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_2 = mma_tile_cursor + tile_2; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n = 0; n < 8; n++) { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n == 0) ? 0 : 1))); + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_1 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + if (tile_2 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int p_stage_1 = prev_phase * 4 + ps_1; + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int output_col = 128 + acc_stage_2 * 128; + int _mma_a_lo_2 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_2 = ((((smem_v_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_2 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_2 = 0; ps_2 < 4; ps_2++) { + int p_stage_last = last_phase * 4 + ps_2; + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int output_col_d = 128 + acc_stage_3 * 128; + int _mma_a_lo_3 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_3 = ((((smem_v_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_2 == 0) ? first_pv : 0)) ? 0 : 1))); + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_2; + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy_1 = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + unsigned int load_kv_stage = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { + int batch_idx_3 = work_idx_3 / (unsigned int)num_split; + int split_idx_3 = work_idx_3 % (unsigned int)num_split; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int tiles_per_split_3 = (total_kv_tiles_3 + num_split - 1) / num_split; + int start_tile_3 = split_idx_3 * tiles_per_split_3; + int end_tile_raw_3 = start_tile_3 + tiles_per_split_3; + int end_tile_3 = ((end_tile_raw_3 < total_kv_tiles_3) ? end_tile_raw_3 : total_kv_tiles_3); + int num_kv_tiles_3 = end_tile_3 - start_tile_3; + int pt_base = batch_idx_3 * max_pages_per_seq; + int q_row_global = batch_idx_3 * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile_3 + 2 * cta_rank; + int pg_k00 = page_table[k0_page_base]; + int pg_k01 = page_table[k0_page_base + 1]; + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + int dst = smem_kv_addr + load_kv_stage * 8192; + tma_4d_gmem2smem_cta2(dst, (&tmap_k_nope), 0, 0, n_1, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k_nope), 0, 0, n_1, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + int dst_r = smem_kv_addr + load_kv_stage * 8192; + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k_rope), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k_rope), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { + int abs_tile_k = start_tile_3 + tile_3; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = page_table[k_page_base]; + int pg_k1 = page_table[k_page_base + 1]; + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + int dst_k = smem_kv_addr + load_kv_stage * 8192; + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k_nope), 0, 0, n_2, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k_nope), 0, 0, n_2, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + int dst_kr = smem_kv_addr + load_kv_stage * 8192; + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k_rope), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k_rope), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + int abs_tile_v = start_tile_3 + tile_3 - 1; + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps_3]; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < 2; acc_stage_4++) { + int v_group = acc_stage_4 * 4 + cta_rank * 2; + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + int dst_v = smem_v_addr + load_kv_stage * 8192; + tma_4d_gmem2smem_cta2(dst_v, (&tmap_v), 0, 0, v_group, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_v + 4096, (&tmap_v), 0, 0, v_group + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + int abs_tile_last = start_tile_3 + num_kv_tiles_3 - 1; + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_4]; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < 2; acc_stage_5++) { + int v_group_last = acc_stage_5 * 4 + cta_rank * 2; + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + int dst_vl = smem_v_addr + load_kv_stage * 8192; + tma_4d_gmem2smem_cta2(dst_vl, (&tmap_v), 0, 0, v_group_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vl + 4096, (&tmap_v), 0, 0, v_group_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + } + // ---- Role: empty0 ---- + if (warp == 10) { + { // empty0_main + const int wg2_dummy_2 = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + #pragma unroll 1 + for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + // ---- Role: empty1 ---- + if (warp == 11) { + { // empty1_main + const int wg2_dummy_3 = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + #pragma unroll 1 + for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_v32_reduce_1db619cce9.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_v32_reduce_1db619cce9.cu new file mode 100644 index 00000000000..38273187961 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_v32_reduce_1db619cce9.cu @@ -0,0 +1,741 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_SPLIT_WEIGHTS_OFF 0 +#define SMEM_SPLIT_WEIGHTS_STAGE_BYTES 128 +#define SMEM_SPLIT_WEIGHTS_STRIDE 128 +#define SMEM_MERGED_LSE_OFF 128 +#define SMEM_MERGED_LSE_STAGE_BYTES 4 +#define SMEM_MERGED_LSE_STRIDE 4 +#define SMEM_TOTAL 256 +#define THREADS 128 +#define USE_PDL 1 + +#include + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_mla_decode_v32_reduce(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, float* __restrict__ sinks, int num_q_heads, int num_split, float bmm2_scale, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + float* split_weights = reinterpret_cast(smem_raw + 0); + const int split_weights_addr = smem + 0; + float* merged_lse = reinterpret_cast(smem_raw + 128); + const int merged_lse_addr = smem + 128; + + // === Task calls (dependency order) === + const int wg_dummy = 0; + int batch_idx = blockIdx.x; + int head_idx = blockIdx.y; + int stat_base = (batch_idx * num_q_heads + head_idx) * num_split; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + if (warp == 0) { + float local_m[1]; + float local_w[1]; + float thread_max = -MLA_INF; + #pragma unroll + for (int i = 0; i < 1; i++) { + int split_idx = lane + (unsigned int)(i * 32); + local_m[i] = -MLA_INF; + local_w[i] = 0.0f; + if (split_idx < num_split) { + local_m[i] = partial_lse[stat_base + split_idx]; + float _max_0 = max_noftz(thread_max, local_m[i]); + thread_max = _max_0; + } + } + float _warp_reduce_0 = thread_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + float global_max = _warp_reduce_0; + float thread_sum = 0.0f; + #pragma unroll + for (int i_1 = 0; i_1 < 1; i_1++) { + int split_idx_1 = lane + (unsigned int)(i_1 * 32); + if (split_idx_1 < num_split) { + float _exp2_0 = approx_exp2(local_m[i_1] - global_max); + local_w[i_1] = ((local_m[i_1] == -MLA_INF) ? 0.0f : _exp2_0); + thread_sum = thread_sum + local_w[i_1]; + } + } + float _warp_reduce_1 = thread_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + float global_sum = _warp_reduce_1; + float _rcp_0 = approx_rcp(global_sum); + float inv_sum = ((global_sum > 0.0f) ? _rcp_0 : 0.0f); + if (lane == 0) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(global_sum)); + merged_lse[0] = global_max + _log2_0; + } + #pragma unroll + for (int i_2 = 0; i_2 < 1; i_2++) { + int split_idx_2 = lane + (unsigned int)(i_2 * 32); + if (split_idx_2 < num_split) { + split_weights[split_idx_2] = local_w[i_2] * inv_sum; + } + } + } + __syncthreads(); + float output_scale = bmm2_scale; + if (enable_sink != 0) { + float _exp2_1 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - merged_lse[0]); + float sink_exp = _exp2_1; + float _rcp_1 = approx_rcp(1.0f + sink_exp); + output_scale = output_scale * _rcp_1; + } + int po_head_base = stat_base * 512; + int o_head_base = (batch_idx * num_q_heads + head_idx) * 512; + float acc[4]; + int d_base = tid * 4; + #pragma unroll + for (int e = 0; e < 4; e++) { + acc[e] = 0.0f; + } + #pragma unroll 4 + for (int s = 0; s < num_split; s++) { + float split_weight = split_weights[s]; + float _vec_load_0[4]; + { + uint2 _vld_0; + _vld_0 = *reinterpret_cast(partial_O + (po_head_base + s * 512 + d_base) + 0); + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&_vec_load_0[0 + _pair * 2])[0]), "=f"((&_vec_load_0[0 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + #pragma unroll + for (int e_1 = 0; e_1 < 4; e_1++) { + acc[e_1] = acc[e_1] + split_weight * _vec_load_0[e_1]; + } + } + { + const float2 _prescale2_1 = {output_scale, output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 2; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&acc[0])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + acc[0 + _ps] *= output_scale; + #endif + uint2 _pk2; + __nv_bfloat162* _pk = reinterpret_cast<__nv_bfloat162*>(&_pk2); + _pk[0] = __floats2bfloat162_rn(acc[0 + 0], acc[0 + 1]); + _pk[1] = __floats2bfloat162_rn(acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (o_head_base + d_base)))[0]) = _pk2; + } + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_value_split_paged_bf16_5c36fe8228.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_value_split_paged_bf16_5c36fe8228.cu new file mode 100644 index 00000000000..a8b1a1b5a8b --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_value_split_paged_bf16_5c36fe8228.cu @@ -0,0 +1,2186 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 74752 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 74752 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 74752 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 197632 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 230400 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 1024 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 1024 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 231424 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_TOTAL 231936 +#define THREADS 384 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384) __cluster_dims__(2,1,1) void +kernel_mla_decode_value_split_paged_bf16(MlaTensorMap const* tmap_q, MlaTensorMap const* tmap_k_nope, MlaTensorMap const* tmap_k_rope, MlaTensorMap const* tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (tid == 0) { + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k_nope)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k_rope)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_v)) : "memory"); + } + __syncthreads(); + + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_kv_addr = smem + 74752; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_vhalf_addr = smem + 74752; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_vquarter_addr = smem + 74752; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 197632); + const int smem_p_addr = smem + 197632; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 230400); + const int smem_softmax_exchange_addr = smem + 230400; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 231424); + const int smem_epilogue_exchange_addr = smem + 231424; + + // Mbarrier init (14 pipeline groups, 0 ordered-sequence groups, 48 barriers) + // Mbarriers at smem_raw[0..384) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + mbarrier_init(smem + 112, 1); + mbarrier_init(smem + 120, 1); + mbarrier_init(smem + 128, 1); + // kv_empty: 15 barriers, init_count=1 + mbarrier_init(smem + 136, 1); + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + mbarrier_init(smem + 160, 1); + mbarrier_init(smem + 168, 1); + mbarrier_init(smem + 176, 1); + mbarrier_init(smem + 184, 1); + mbarrier_init(smem + 192, 1); + mbarrier_init(smem + 200, 1); + mbarrier_init(smem + 208, 1); + mbarrier_init(smem + 216, 1); + mbarrier_init(smem + 224, 1); + mbarrier_init(smem + 232, 1); + mbarrier_init(smem + 240, 1); + mbarrier_init(smem + 248, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 256, 1); + mbarrier_init(smem + 264, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 272, 256); + mbarrier_init(smem + 280, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 288, 256); + mbarrier_init(smem + 296, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 304, 1); + mbarrier_init(smem + 312, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 320, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 328, 128); + mbarrier_init(smem + 336, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 344, 128); + mbarrier_init(smem + 352, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 360, 1); + // tmem_scrubbed: 1 barriers, init_count=256 + mbarrier_init(smem + 368, 256); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 376, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 384); + if (warp == 0) { + int _tmem_hold = smem + 384; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + #pragma unroll 1 + for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { + int batch_idx = work_idx; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + #pragma unroll 1 + for (int tile = 0; tile < num_kv_tiles; tile++) { + int pipeline_tile = softmax_tile_cursor + tile; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile + tile; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[phase * 128 + exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[phase * 128 + exchange_idx ^ 64]); + new_max = _max_1; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (phase) * 8, p_cor_empty_phase); + float meta[4]; + meta[0] = row_sum_val; + meta[1] = row_max_val; + meta[2] = acc_scale; + meta[3] = no_correction; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + tmem_st_x4_f32(meta_addr, meta); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { + int batch_idx_1 = work_idx_1; + int value_split = 0; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + if (num_kv_tiles_1 > 0) { + int first_phase = correction_tile_cursor & 1; + int first_wait_phase = correction_tile_cursor >> 1 & 1; + mbarrier_wait(stats_addr + (first_phase) * 8, first_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int first_meta_addr = taddr + 384 + (unsigned int)(first_phase * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], first_meta_addr); + final_sum_val = _tmem_load_0[0]; + final_max_val = _tmem_load_0[1]; + mbarrier_arrive(stats_empty_addr + (first_phase) * 8); + } + #pragma unroll 1 + for (int tile_1 = 1; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile_1 = correction_tile_cursor + tile_1; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], meta_addr_1); + final_sum_val = _tmem_load_1[0]; + final_max_val = _tmem_load_1[1]; + float acc_scale_1 = _tmem_load_1[2]; + float no_correction_1 = _tmem_load_1[3]; + { + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); + int skip_correction = _vote_0; + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage = 0; acc_stage < 2; acc_stage++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_3[16]; + tmem_ld_x16(&_tmem_load_3[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_3)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_3)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_3); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_1 * 128 + head_idx) * 512; + float output_scale = bmm2_scale; + if (enable_sink != 0) { + float _exp2_3 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - lse_log2); + float sink_exp = _exp2_3; + float _rcp_1 = approx_rcp(1.0f + sink_exp); + output_scale = output_scale * _rcp_1; + } + { + { + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_1 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + const float2 _scale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_8)[_ls] = mul_f32x2_rn_noftz(reinterpret_cast(_tmem_load_8)[_ls], _scale2_1); + unsigned int pair_packed[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_8[_lp*2 + 0], _tmem_load_8[_lp*2+1 + 0])); + pair_packed[_lp] = *(uint32_t*)&_bf2; + } + int pair_parity = lane & 1; + unsigned int pair_mask = (unsigned int)(-pair_parity); + int pair_neighbor = lane ^ 1; + int pair_even_base = gmem_base - pair_parity * 512 + pair_parity * 8; + #pragma unroll + for (int pair_columns = 0; pair_columns < 2; pair_columns++) { + unsigned int pair_even_words[4]; + unsigned int pair_odd_words[4]; + #pragma unroll + for (int pair_word = 0; pair_word < 4; pair_word++) { + unsigned int _shfl_0 = __shfl_sync(0xFFFFFFFF, pair_packed[pair_columns * 8 + 4 + pair_word] & ~pair_mask | pair_packed[pair_columns * 8 + pair_word] & pair_mask, pair_neighbor); + pair_even_words[pair_word] = pair_packed[pair_columns * 8 + pair_word] & ~pair_mask | _shfl_0 & pair_mask; + pair_odd_words[pair_word] = _shfl_0 & ~pair_mask | pair_packed[pair_columns * 8 + 4 + pair_word] & pair_mask; + } + reinterpret_cast(O + (pair_even_base + pair_columns * 16))[0] = reinterpret_cast(pair_even_words)[0]; + reinterpret_cast(O + (pair_even_base + 512 + pair_columns * 16))[0] = reinterpret_cast(pair_odd_words)[0]; + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_1; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_k_nope)) : "memory"); + { + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_k_rope)) : "memory"); + } + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_v)) : "memory"); + const int wg2_dummy = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + #pragma unroll 1 + for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { + int batch_idx_2 = work_idx_2; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_2 = mma_tile_cursor + tile_2; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n = 0; n < 8; n++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_2 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int p_stage_1 = prev_phase * 4 + ps_1; + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < ((1) ? 2 : 1); acc_stage_2++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_2 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_2 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_2 = 0; ps_2 < 4; ps_2++) { + int p_stage_last = last_phase * 4 + ps_2; + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < ((1) ? 2 : 1); acc_stage_3++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_3 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_2 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + { + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy_1 = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { + int batch_idx_3 = work_idx_3; + int value_split_1 = 0; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + int pt_base = batch_idx_3 * max_pages_per_seq; + int q_row_global = batch_idx_3 * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, tmap_q, 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile_3 + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, tmap_k_nope, 0, 0, n_1, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, tmap_k_nope, 0, 0, n_1, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, tmap_k_rope, 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, tmap_k_rope, 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { + int abs_tile_k = start_tile_3 + tile_3; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, tmap_k_nope, 0, 0, n_2, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, tmap_k_nope, 0, 0, n_2, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, tmap_k_rope, 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, tmap_k_rope, 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile_3 + tile_3 - 1; + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps_3]; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage_4 * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, tmap_v, 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, tmap_v, 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile_3 + num_kv_tiles_3 - 1; + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_4]; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_5 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, tmap_v, 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, tmap_v, 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + } + // ---- Role: empty0 ---- + if (warp == 10) { + { // empty0_main + const int wg2_dummy_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + // ---- Role: empty1 ---- + if (warp == 11) { + { // empty1_main + const int wg2_dummy_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_value_split_paged_bf16_ee41d3ceea.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_value_split_paged_bf16_ee41d3ceea.cu new file mode 100644 index 00000000000..06289165c0d --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_value_split_paged_bf16_ee41d3ceea.cu @@ -0,0 +1,2131 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 74752 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 74752 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 74752 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 197632 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 230400 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 230912 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_TOTAL 231424 +#define THREADS 384 +#define V_QUARTER 1 +#define UNSPLIT 0 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384) __cluster_dims__(2,1,1) void +kernel_mla_decode_value_split_paged_bf16(MlaTensorMap const* tmap_q, MlaTensorMap const* tmap_k_nope, MlaTensorMap const* tmap_k_rope, MlaTensorMap const* tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (tid == 0) { + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k_nope)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k_rope)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_v)) : "memory"); + } + __syncthreads(); + + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_kv_addr = smem + 74752; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_vhalf_addr = smem + 74752; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_vquarter_addr = smem + 74752; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 197632); + const int smem_p_addr = smem + 197632; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 230400); + const int smem_softmax_exchange_addr = smem + 230400; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 230912); + const int smem_epilogue_exchange_addr = smem + 230912; + + // Mbarrier init (14 pipeline groups, 0 ordered-sequence groups, 48 barriers) + // Mbarriers at smem_raw[0..384) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + mbarrier_init(smem + 112, 1); + mbarrier_init(smem + 120, 1); + mbarrier_init(smem + 128, 1); + // kv_empty: 15 barriers, init_count=1 + mbarrier_init(smem + 136, 1); + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + mbarrier_init(smem + 160, 1); + mbarrier_init(smem + 168, 1); + mbarrier_init(smem + 176, 1); + mbarrier_init(smem + 184, 1); + mbarrier_init(smem + 192, 1); + mbarrier_init(smem + 200, 1); + mbarrier_init(smem + 208, 1); + mbarrier_init(smem + 216, 1); + mbarrier_init(smem + 224, 1); + mbarrier_init(smem + 232, 1); + mbarrier_init(smem + 240, 1); + mbarrier_init(smem + 248, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 256, 1); + mbarrier_init(smem + 264, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 272, 256); + mbarrier_init(smem + 280, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 288, 256); + mbarrier_init(smem + 296, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 304, 1); + mbarrier_init(smem + 312, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 320, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 328, 128); + mbarrier_init(smem + 336, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 344, 128); + mbarrier_init(smem + 352, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 360, 1); + // tmem_scrubbed: 1 barriers, init_count=256 + mbarrier_init(smem + 368, 256); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 376, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 384); + if (warp == 0) { + int _tmem_hold = smem + 384; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + #pragma unroll 1 + for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { + int batch_idx = work_idx / (unsigned int)value_split_count; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + #pragma unroll 1 + for (int tile = 0; tile < num_kv_tiles; tile++) { + int pipeline_tile = softmax_tile_cursor + tile; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile + tile; + bool score_tile_full = 0; + int tail_valid = seqlen_kv_b - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + #pragma unroll + for (int i = 0; i < 64; i++) { + float _exp2_1 = approx_exp2(sv[i] * softmax_scale_log2 - max_scaled); + sv[i] = _exp2_1; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (phase) * 8, p_cor_empty_phase); + float meta[4]; + meta[0] = row_sum_val; + meta[1] = row_max_val; + meta[2] = acc_scale; + meta[3] = no_correction; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + tmem_st_x4_f32(meta_addr, meta); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { + int batch_idx_1 = work_idx_1 / (unsigned int)value_split_count; + int value_split = work_idx_1 % (unsigned int)value_split_count; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile_1 = correction_tile_cursor + tile_1; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], meta_addr_1); + final_sum_val = _tmem_load_0[0]; + final_max_val = _tmem_load_0[1]; + float acc_scale_1 = _tmem_load_0[2]; + float no_correction_1 = _tmem_load_0[3]; + if (tile_1 > 0) { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + if (tile_1 > 0) { + int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + #pragma unroll + for (int vs_local = 0; vs_local < 1; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_1[16]; + tmem_ld_x16(&_tmem_load_1[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_1)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_1)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_1); + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_1 * 128 + head_idx) * 512; + float output_scale = bmm2_scale; + if (enable_sink != 0) { + float _exp2_2 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - lse_log2); + float sink_exp = _exp2_2; + float _rcp_1 = approx_rcp(1.0f + sink_exp); + output_scale = output_scale * _rcp_1; + } + { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 1; vs_local_1++) { + int logical_vs = value_split * 2 + n_half_1 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_4[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]), "=f"(_tmem_load_4[2]), "=f"(_tmem_load_4[3]), "=f"(_tmem_load_4[4]), "=f"(_tmem_load_4[5]), "=f"(_tmem_load_4[6]), "=f"(_tmem_load_4[7]), "=f"(_tmem_load_4[8]), "=f"(_tmem_load_4[9]), "=f"(_tmem_load_4[10]), "=f"(_tmem_load_4[11]), "=f"(_tmem_load_4[12]), "=f"(_tmem_load_4[13]), "=f"(_tmem_load_4[14]), "=f"(_tmem_load_4[15]), "=f"(_tmem_load_4[16]), "=f"(_tmem_load_4[17]), "=f"(_tmem_load_4[18]), "=f"(_tmem_load_4[19]), "=f"(_tmem_load_4[20]), "=f"(_tmem_load_4[21]), "=f"(_tmem_load_4[22]), "=f"(_tmem_load_4[23]), "=f"(_tmem_load_4[24]), "=f"(_tmem_load_4[25]), "=f"(_tmem_load_4[26]), "=f"(_tmem_load_4[27]), "=f"(_tmem_load_4[28]), "=f"(_tmem_load_4[29]), "=f"(_tmem_load_4[30]), "=f"(_tmem_load_4[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_4[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_4[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_4[j + 0], _tmem_load_4[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_4[j + 2], _tmem_load_4[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_4[j + 4], _tmem_load_4[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_4[j + 6], _tmem_load_4[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_1; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_k_nope)) : "memory"); + { + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_k_rope)) : "memory"); + } + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_v)) : "memory"); + const int wg2_dummy = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + #pragma unroll 1 + for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { + int batch_idx_2 = work_idx_2 / (unsigned int)value_split_count; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_2 = mma_tile_cursor + tile_2; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n = 0; n < 8; n++) { + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (n) * 512; + int _mma_b_lo_1 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(((n == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int _mma_a_lo_3 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_3 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + if (tile_2 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int p_stage_1 = prev_phase * 4 + ps_1; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((UNSPLIT) ? 2 : 1); acc_stage++) { + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int output_col = 128 + acc_stage * 128; + { + int _mma_a_lo_5 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_5 = ((((smem_vquarter_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136381584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_2 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_2 = 0; ps_2 < 4; ps_2++) { + int p_stage_last = last_phase * 4 + ps_2; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((UNSPLIT) ? 2 : 1); acc_stage_1++) { + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int output_col_d = 128 + acc_stage_1 * 128; + { + int _mma_a_lo_8 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_8 = ((((smem_vquarter_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136381584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_8), "r"(_mma_b_lo_8), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_2 == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + { + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy_1 = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { + int batch_idx_3 = work_idx_3 / (unsigned int)value_split_count; + int value_split_1 = work_idx_3 % (unsigned int)value_split_count; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + int pt_base = batch_idx_3 * max_pages_per_seq; + int q_row_global = batch_idx_3 * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, tmap_q, 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile_3 + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst, tmap_k_nope, 0, 0, n_1, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, tmap_k_nope, 0, 0, n_1, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + int dst_r = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_r = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_r, tmap_k_rope, 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, tmap_k_rope, 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { + int abs_tile_k = start_tile_3 + tile_3; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + int dst_k = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_k = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_k, tmap_k_nope, 0, 0, n_2, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, tmap_k_nope, 0, 0, n_2, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + int dst_kr = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_kr = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_kr, tmap_k_rope, 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, tmap_k_rope, 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + int abs_tile_v = start_tile_3 + tile_3 - 1; + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps_3]; + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < ((UNSPLIT) ? 2 : 1); acc_stage_2++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + { + int v_group_quarter = value_split_1 * 2 + cta_rank; + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(8192)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(smem_vquarter_addr + load_kv_stage * 8192, tmap_v, 0, 0, v_group_quarter, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + } + int abs_tile_last = start_tile_3 + num_kv_tiles_3 - 1; + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_4]; + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < ((UNSPLIT) ? 2 : 1); acc_stage_3++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + { + int v_group_quarter_last = value_split_1 * 2 + cta_rank; + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(8192)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(smem_vquarter_addr + load_kv_stage * 8192, tmap_v, 0, 0, v_group_quarter_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + } + // ---- Role: empty0 ---- + if (warp == 10) { + { // empty0_main + const int wg2_dummy_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + // ---- Role: empty1 ---- + if (warp == 11) { + { // empty1_main + const int wg2_dummy_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_value_split_paged_bf16_fa7b59f3a7.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_value_split_paged_bf16_fa7b59f3a7.cu new file mode 100644 index 00000000000..ae89d8a271d --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_mla_decode_value_split_paged_bf16_fa7b59f3a7.cu @@ -0,0 +1,2147 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 74752 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 74752 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 74752 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 197632 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 230400 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 230912 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_TOTAL 231424 +#define THREADS 384 +#define V_QUARTER 0 +#define UNSPLIT 0 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384) __cluster_dims__(2,1,1) void +kernel_mla_decode_value_split_paged_bf16(MlaTensorMap const* tmap_q, MlaTensorMap const* tmap_k_nope, MlaTensorMap const* tmap_k_rope, MlaTensorMap const* tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (tid == 0) { + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k_nope)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k_rope)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_v)) : "memory"); + } + __syncthreads(); + + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_kv_addr = smem + 74752; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_vhalf_addr = smem + 74752; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_vquarter_addr = smem + 74752; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 197632); + const int smem_p_addr = smem + 197632; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 230400); + const int smem_softmax_exchange_addr = smem + 230400; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 230912); + const int smem_epilogue_exchange_addr = smem + 230912; + + // Mbarrier init (14 pipeline groups, 0 ordered-sequence groups, 48 barriers) + // Mbarriers at smem_raw[0..384) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + mbarrier_init(smem + 112, 1); + mbarrier_init(smem + 120, 1); + mbarrier_init(smem + 128, 1); + // kv_empty: 15 barriers, init_count=1 + mbarrier_init(smem + 136, 1); + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + mbarrier_init(smem + 160, 1); + mbarrier_init(smem + 168, 1); + mbarrier_init(smem + 176, 1); + mbarrier_init(smem + 184, 1); + mbarrier_init(smem + 192, 1); + mbarrier_init(smem + 200, 1); + mbarrier_init(smem + 208, 1); + mbarrier_init(smem + 216, 1); + mbarrier_init(smem + 224, 1); + mbarrier_init(smem + 232, 1); + mbarrier_init(smem + 240, 1); + mbarrier_init(smem + 248, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 256, 1); + mbarrier_init(smem + 264, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 272, 256); + mbarrier_init(smem + 280, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 288, 256); + mbarrier_init(smem + 296, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 304, 1); + mbarrier_init(smem + 312, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 320, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 328, 128); + mbarrier_init(smem + 336, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 344, 128); + mbarrier_init(smem + 352, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 360, 1); + // tmem_scrubbed: 1 barriers, init_count=256 + mbarrier_init(smem + 368, 256); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 376, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 384); + if (warp == 0) { + int _tmem_hold = smem + 384; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + #pragma unroll 1 + for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { + int batch_idx = work_idx / (unsigned int)value_split_count; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + #pragma unroll 1 + for (int tile = 0; tile < num_kv_tiles; tile++) { + int pipeline_tile = softmax_tile_cursor + tile; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile + tile; + bool score_tile_full = 0; + int tail_valid = seqlen_kv_b - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + #pragma unroll + for (int i = 0; i < 64; i++) { + float _exp2_1 = approx_exp2(sv[i] * softmax_scale_log2 - max_scaled); + sv[i] = _exp2_1; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (phase) * 8, p_cor_empty_phase); + float meta[4]; + meta[0] = row_sum_val; + meta[1] = row_max_val; + meta[2] = acc_scale; + meta[3] = no_correction; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + tmem_st_x4_f32(meta_addr, meta); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { + int batch_idx_1 = work_idx_1 / (unsigned int)value_split_count; + int value_split = work_idx_1 % (unsigned int)value_split_count; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile_1 = correction_tile_cursor + tile_1; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], meta_addr_1); + final_sum_val = _tmem_load_0[0]; + final_max_val = _tmem_load_0[1]; + float acc_scale_1 = _tmem_load_0[2]; + float no_correction_1 = _tmem_load_0[3]; + if (tile_1 > 0) { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + if (tile_1 > 0) { + int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_3[16]; + tmem_ld_x16(&_tmem_load_3[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_3)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_3)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_3); + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_1 * 128 + head_idx) * 512; + float output_scale = bmm2_scale; + if (enable_sink != 0) { + float _exp2_2 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - lse_log2); + float sink_exp = _exp2_2; + float _rcp_1 = approx_rcp(1.0f + sink_exp); + output_scale = output_scale * _rcp_1; + } + { + { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = value_split * 4 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_8[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_8[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_8[j + 0], _tmem_load_8[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_8[j + 2], _tmem_load_8[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_8[j + 4], _tmem_load_8[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_8[j + 6], _tmem_load_8[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_1; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_k_nope)) : "memory"); + { + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_k_rope)) : "memory"); + } + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_v)) : "memory"); + const int wg2_dummy = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + #pragma unroll 1 + for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { + int batch_idx_2 = work_idx_2 / (unsigned int)value_split_count; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_2 = mma_tile_cursor + tile_2; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n = 0; n < 8; n++) { + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (n) * 512; + int _mma_b_lo_1 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(((n == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int _mma_a_lo_3 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_3 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + if (tile_2 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int p_stage_1 = prev_phase * 4 + ps_1; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((UNSPLIT) ? 2 : 1); acc_stage++) { + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int output_col = 128 + acc_stage * 128; + { + int _mma_a_lo_6 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_6 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_2 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_2 = 0; ps_2 < 4; ps_2++) { + int p_stage_last = last_phase * 4 + ps_2; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((UNSPLIT) ? 2 : 1); acc_stage_1++) { + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int output_col_d = 128 + acc_stage_1 * 128; + { + int _mma_a_lo_9 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_9 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_9), "r"(_mma_b_lo_9), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_2 == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + { + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy_1 = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { + int batch_idx_3 = work_idx_3 / (unsigned int)value_split_count; + int value_split_1 = work_idx_3 % (unsigned int)value_split_count; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + int pt_base = batch_idx_3 * max_pages_per_seq; + int q_row_global = batch_idx_3 * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, tmap_q, 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile_3 + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst, tmap_k_nope, 0, 0, n_1, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, tmap_k_nope, 0, 0, n_1, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + int dst_r = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_r = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_r, tmap_k_rope, 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, tmap_k_rope, 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { + int abs_tile_k = start_tile_3 + tile_3; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + int dst_k = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_k = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_k, tmap_k_nope, 0, 0, n_2, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, tmap_k_nope, 0, 0, n_2, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + int dst_kr = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_kr = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_kr, tmap_k_rope, 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, tmap_k_rope, 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + int abs_tile_v = start_tile_3 + tile_3 - 1; + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps_3]; + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < ((UNSPLIT) ? 2 : 1); acc_stage_2++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + { + { + v_group_half = value_split_1 * 4 + cta_rank * 2; + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_vhalf, tmap_v, 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, tmap_v, 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + } + int abs_tile_last = start_tile_3 + num_kv_tiles_3 - 1; + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_4]; + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < ((UNSPLIT) ? 2 : 1); acc_stage_3++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + { + { + v_group_half_last = value_split_1 * 4 + cta_rank * 2; + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_vhalf_last, tmap_v, 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, tmap_v, 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + } + // ---- Role: empty0 ---- + if (warp == 10) { + { // empty0_main + const int wg2_dummy_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + // ---- Role: empty1 ---- + if (warp == 11) { + { // empty1_main + const int wg2_dummy_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_page64_keeps_gmem_separate_reduce_cea01a87e9.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_page64_keeps_gmem_separate_reduce_cea01a87e9.cu new file mode 100644 index 00000000000..667f682d4b4 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_page64_keeps_gmem_separate_reduce_cea01a87e9.cu @@ -0,0 +1,704 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 512 +#define USE_PDL 1 + +#include + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +extern "C" { + +__global__ __launch_bounds__(512) void +kernel_page64_keeps_gmem_separate_reduce(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_stats, __nv_bfloat16* __restrict__ O, float* __restrict__ LSE, int num_split, int num_reduce_ctas, int num_queries, float bmm2_scale) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + const int role_token = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + int reduction_cta = blockIdx.x; + int work_item = blockIdx.z; + int partition = blockIdx.y; + int head_group = partition; + int v_partition = 0; + int slices_per_cta = (8 + num_reduce_ctas - 1) / num_reduce_ctas; + int slice_begin = reduction_cta * slices_per_cta; + int slice_end_raw = slice_begin + slices_per_cta; + int slice_end = ((slice_end_raw < 8) ? slice_end_raw : 8); + #pragma unroll 1 + for (int slice_idx = slice_begin; slice_idx < slice_end; slice_idx++) { + int vector_idx = tid; + int row_in_group = slice_idx * 8 + vector_idx / 64; + int dim = v_partition * 512 + vector_idx % 64 * 8; + int head = head_group * 64 + row_in_group; + float merged[8]; + #pragma unroll + for (int item = 0; item < 8; item++) { + merged[item] = 0.0f; + } + float merged_max = -MLA_INF; + float merged_sum = 0.0f; + #pragma unroll 4 + for (int split = 0; split < num_split; split++) { + int stat_base = ((work_item * num_split + split) * 128 + head) * 2; + float local_max = partial_stats[stat_base]; + float local_sum = partial_stats[stat_base + 1]; + float _max_0 = max_noftz(merged_max, local_max); + float next_max = _max_0; + float _exp2_0 = approx_exp2(merged_max - next_max); + float prior_scale = ((merged_max != -MLA_INF) ? _exp2_0 : 0.0f); + float _exp2_1 = approx_exp2(local_max - next_max); + float local_scale = ((local_max != -MLA_INF) ? _exp2_1 : 0.0f); + int partial_base = ((work_item * num_split + split) * 128 + head) * 512 + dim; + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(partial_O + partial_base + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&_vec_load_0[0 + _blk * 8 + _pair * 2])[0]), "=f"((&_vec_load_0[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + #pragma unroll + for (int item_1 = 0; item_1 < 8; item_1++) { + merged[item_1] = merged[item_1] * prior_scale + _vec_load_0[item_1] * local_scale; + } + merged_sum = merged_sum * prior_scale + local_sum * local_scale; + merged_max = next_max; + } + float _rcp_0 = approx_rcp(merged_sum); + float inv_sum = ((merged_sum > 0.0f) ? _rcp_0 : 0.0f); + int output_base = (work_item * 128 + head) * 512 + dim; + { + const float2 _prescale2_1 = {inv_sum * bmm2_scale, inv_sum * bmm2_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&merged[0])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + merged[0 + _ps] *= inv_sum * bmm2_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(merged[0 + 0], merged[0 + 1]); + _pk[1] = __floats2bfloat162_rn(merged[0 + 2], merged[0 + 3]); + _pk[2] = __floats2bfloat162_rn(merged[0 + 4], merged[0 + 5]); + _pk[3] = __floats2bfloat162_rn(merged[0 + 6], merged[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + output_base))[0]) = *reinterpret_cast(&_pk[0]); + } + if (dim == 0) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(merged_sum)); + LSE[work_item * 128 + head] = ((merged_sum > 0.0f) ? merged_max + _log2_0 : -MLA_INF); + } + } + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_trtllm_mla_decode_fp8_06282d4d84.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_trtllm_mla_decode_fp8_06282d4d84.cu new file mode 100644 index 00000000000..b75fd8d3410 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_100a_148/cake_trtllm_mla_decode_fp8_06282d4d84.cu @@ -0,0 +1,3271 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_K_PIPE_STAGES 2 +#define NUM_V_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 40960 +#define SMEM_SMEM_Q_STRIDE 40960 +#define SMEM_SMEM_K_OFF 41984 +#define SMEM_SMEM_K_STAGE_BYTES 40960 +#define SMEM_SMEM_K_STRIDE 40960 +#define SMEM_SMEM_V_OFF 123904 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_SMEM_STATS_MAX_OFF 189440 +#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 2048 +#define SMEM_SMEM_STATS_MAX_STRIDE 2048 +#define SMEM_SMEM_STATS_SUM_OFF 191488 +#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 1280 +#define SMEM_SMEM_STATS_SUM_STRIDE 1280 +#define SMEM_SMEM_P_OFF 193536 +#define SMEM_SMEM_P_STAGE_BYTES 8192 +#define SMEM_SMEM_P_STRIDE 8192 +#define SMEM_SMEM_EPI_OFF 209920 +#define SMEM_SMEM_EPI_STAGE_BYTES 16384 +#define SMEM_SMEM_EPI_STRIDE 16384 +#define SMEM_TOTAL 226304 +#define THREADS 384 +#define USE_PDL 1 +#define SM103_SCORE_LD_RED_MAX 0 +#define SM100_SCORE_TREE_MAX 1 +#define SM100_ROW_SUM_CHAINS 1 +#define SM100_DEFER_ROW_SUM 0 +#define SM100_CORRECTION_BATCH 1 +#define SM100_F16X2_EXP 0 +#define SM100_EXP2_EMU25 1 +#define SM100_EXP2_EMU44 0 +#define SM100_EXP2_EMU62 0 +#define SM100_DROP_PROXY_FENCES 1 +#define SM100_EXP2_EMU31 0 +#define SM100_EXP2_EMU41 0 +#define SM100_PACKED_MATH 1 +#define SM100_KO_NO_EXP 0 +#define SM100_KO_NO_RESCALE 0 +#define SM100_KO_NO_EXCHANGE 0 +#define SM100_KO_NO_P_STORE 0 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ uint32_t mbarrier_test_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.test_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_test_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +__device__ __forceinline__ void mbarrier_wait_token_test(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_TEST:\n\t" + "mbarrier.test_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_TEST;\n\t" + "bra.uni LAB_WAIT_TEST;\n\t" + "DONE_TEST:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_test_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_TEST_CLUSTER:\n\t" + "mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_TEST_CLUSTER;\n\t" + "bra.uni LAB_WAIT_TEST_CLUSTER;\n\t" + "DONE_TEST_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x32_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x32.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16," + " %17, %18, %19, %20, %21, %22, %23, %24," + " %25, %26, %27, %28, %29, %30, %31, %32};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15]), + "f"(src[16]), "f"(src[17]), "f"(src[18]), "f"(src[19]), + "f"(src[20]), "f"(src[21]), "f"(src[22]), "f"(src[23]), + "f"(src[24]), "f"(src[25]), "f"(src[26]), "f"(src[27]), + "f"(src[28]), "f"(src[29]), "f"(src[30]), "f"(src[31])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" + :: "r"(addr), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), + "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_trtllm_mla_decode_fp8(MlaTensorMap const* tmap_q, MlaTensorMap const* tmap_k, MlaTensorMap const* tmap_v, __nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_stats, __nv_bfloat16* __restrict__ O, float* __restrict__ LSE, int* __restrict__ seq_lens_kv, int* __restrict__ work_batch_indices, int* __restrict__ page_table, float softmax_scale_log2, float bmm2_scale, int num_split, int num_queries, int max_pages_per_seq) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define k_full_addr (mbar_base + 16) + #define k_empty_addr (mbar_base + 32) + #define v_full_addr (mbar_base + 48) + #define v_empty_addr (mbar_base + 64) + #define s_full_addr (mbar_base + 80) + #define s_empty_addr (mbar_base + 96) + #define p_full_addr (mbar_base + 112) + #define p_empty_addr (mbar_base + 128) + #define o_empty_addr (mbar_base + 144) + #define sum_ready_addr (mbar_base + 160) + #define o_full_addr (mbar_base + 168) + #define o_done_addr (mbar_base + 184) + #define partial_results_ready_addr (mbar_base + 192) + #define q_pair_ready_addr (mbar_base + 200) + #define tmem_dealloc_peer_addr (mbar_base + 208) + #define tmem_epoch_ready_addr (mbar_base + 216) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (tid == 0) { + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_v)) : "memory"); + } + __syncthreads(); + + + // Kernel setup ops + uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + uint8_t* smem_k = reinterpret_cast(smem_raw + 41984); + const int smem_k_addr = smem + 41984; + uint8_t* smem_v = reinterpret_cast(smem_raw + 123904); + const int smem_v_addr = smem + 123904; + float* smem_stats_max = reinterpret_cast(smem_raw + 189440); + const int smem_stats_max_addr = smem + 189440; + float* smem_stats_sum = reinterpret_cast(smem_raw + 191488); + const int smem_stats_sum_addr = smem + 191488; + uint8_t* smem_p = reinterpret_cast(smem_raw + 193536); + const int smem_p_addr = smem + 193536; + unsigned int* smem_epi = reinterpret_cast(smem_raw + 209920); + const int smem_epi_addr = smem + 209920; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 28 barriers) + // Mbarriers at smem_raw[0..224) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'k_pipe' --- + // k_full: 2 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + // k_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + // --- pipeline 'v_pipe' --- + // v_full: 2 barriers, init_count=1 + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + // v_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + // s_empty: 2 barriers, init_count=16 + mbarrier_init(smem + 96, 16); + mbarrier_init(smem + 104, 16); + // p_full: 2 barriers, init_count=16 + mbarrier_init(smem + 112, 16); + mbarrier_init(smem + 120, 16); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 128, 1); + mbarrier_init(smem + 136, 1); + // o_empty: 2 barriers, init_count=8 + mbarrier_init(smem + 144, 8); + mbarrier_init(smem + 152, 8); + // sum_ready: 1 barriers, init_count=256 + mbarrier_init(smem + 160, 256); + // o_full: 2 barriers, init_count=1 + mbarrier_init(smem + 168, 1); + mbarrier_init(smem + 176, 1); + // o_done: 1 barriers, init_count=1 + mbarrier_init(smem + 184, 1); + // partial_results_ready: 1 barriers, init_count=256 + mbarrier_init(smem + 192, 256); + // q_pair_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 200, 64); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 208, 32); + // tmem_epoch_ready: 1 barriers, init_count=256 + mbarrier_init(smem + 216, 256); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 224); + if (warp == 0) { + int _tmem_hold = smem + 224; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 112;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 200;"); + { // softmax_wg_main + uint32_t _mbar_token_0 = mbarrier_test_wait(q_full_addr, 0); + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int score_half = warp / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + int split_idx = blockIdx.x / 2; + int batch_idx = blockIdx.y; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int num_kv_tiles_total = (seqlen_kv_b + 128 - 1) / 128; + int blocks_per_split = (num_kv_tiles_total + num_split - 1) / num_split; + int my_start_block = split_idx * blocks_per_split; + int my_end_block_raw = my_start_block + blocks_per_split; + int my_end_block = ((my_end_block_raw < num_kv_tiles_total) ? my_end_block_raw : num_kv_tiles_total); + int my_n_tiles = my_end_block - my_start_block; + mbarrier_wait_token_test(q_full_addr, 0, _mbar_token_0); + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + #pragma unroll 1 + for (int tile = 0; tile < my_n_tiles; tile++) { + int phase = tile & 1; + int s_wait_phase = tile >> 1 & 1; + float sv[32]; + float local_max = -MLA_INF; + { + { + uint32_t _mbar_token_2 = mbarrier_test_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 128 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + int abs_tile = my_start_block + tile; + int tail_valid = seqlen_kv_b - abs_tile * 128; + int local_tail_raw = tail_valid - score_half * 64; + int local_tail_hi = ((local_tail_raw > 64) ? 64 : local_tail_raw); + int local_tail = ((local_tail_hi < 0) ? 0 : local_tail_hi); + mbarrier_wait_token(s_full_addr + (phase) * 8, s_wait_phase, _mbar_token_2); + asm volatile("tcgen05.fence::after_thread_sync;"); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + int sm_tail = ((local_tail > 32) ? 32 : local_tail); + if (sm_tail < 32) { + uint32_t _slice_lo_mask_1; + { + int _lim_0 = sm_tail; + if (_lim_0 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[31] = -MLA_INF; + } + float _max3_0; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_0) : "f"(sv[0]), "f"(sv[1]), "f"(sv[2])); + float _max3_1; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_1) : "f"(sv[3]), "f"(sv[4]), "f"(sv[5])); + float _max3_2; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_2) : "f"(sv[6]), "f"(sv[7]), "f"(sv[8])); + float _max3_3; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_3) : "f"(sv[9]), "f"(sv[10]), "f"(sv[11])); + float _max3_4; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_4) : "f"(sv[12]), "f"(sv[13]), "f"(sv[14])); + float _max3_5; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_5) : "f"(sv[15]), "f"(sv[16]), "f"(sv[17])); + float _max3_6; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_6) : "f"(sv[18]), "f"(sv[19]), "f"(sv[20])); + float _max3_7; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_7) : "f"(sv[21]), "f"(sv[22]), "f"(sv[23])); + float _max3_8; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_8) : "f"(sv[24]), "f"(sv[25]), "f"(sv[26])); + float _max3_9; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_9) : "f"(sv[27]), "f"(sv[28]), "f"(sv[29])); + float _max_1 = max_noftz(sv[30], sv[31]); + float _max3_10; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_10) : "f"(_max3_0), "f"(_max3_1), "f"(_max3_2)); + float _max3_11; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_11) : "f"(_max3_3), "f"(_max3_4), "f"(_max3_5)); + float _max3_12; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_12) : "f"(_max3_6), "f"(_max3_7), "f"(_max3_8)); + float _max_2 = max_noftz(_max3_9, _max_1); + float _max3_13; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_13) : "f"(_max3_10), "f"(_max3_11), "f"(_max3_12)); + float _max_3 = max_noftz(_max3_13, _max_2); + local_max = _max_3; + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + __syncwarp(); + if (elect_sync()) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + int pair_max_idx = phase * 256 + score_half * 128 + my_row; + smem_stats_max[pair_max_idx] = local_max; + { + if (warp % 2 == 0) { + asm volatile("barrier.sync 6, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 7, 128;" ::: "memory"); + } + } + asm volatile("tcgen05.fence::after_thread_sync;"); + float new_max = local_max; + { + float mx_c = smem_stats_max[phase * 256 + score_half * 128 + 64 + my_row]; + float mx_s2 = smem_stats_max[phase * 256 + (1 - score_half) * 128 + my_row]; + float mx_c2 = smem_stats_max[phase * 256 + (1 - score_half) * 128 + 64 + my_row]; + float _max3_14; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_14) : "f"(local_max), "f"(mx_c), "f"(mx_s2)); + new_max = _max3_14; + float _max_5 = max_noftz(new_max, mx_c2); + new_max = _max_5; + } + float _max_6 = max_noftz(new_max, row_max_val); + new_max = _max_6; + float lazy_delta = softmax_scale_log2 * (row_max_val - new_max); + new_max = ((lazy_delta < -8.0f) ? new_max : row_max_val); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + { + { + { + float neg_max_scaled = -max_scaled; + const float2 _fma_b2_1 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_2 = {neg_max_scaled, neg_max_scaled}; + float2 _fma_pair_3 = fma_f32x2(make_float2(sv[0], sv[1]), _fma_b2_1, _fma_c2_2); + sv[0] = _fma_pair_3.x; + sv[1] = _fma_pair_3.y; + float2 _fma_pair_4 = fma_f32x2(make_float2(sv[2], sv[3]), _fma_b2_1, _fma_c2_2); + sv[2] = _fma_pair_4.x; + sv[3] = _fma_pair_4.y; + float2 _fma_pair_5 = fma_f32x2(make_float2(sv[4], sv[5]), _fma_b2_1, _fma_c2_2); + sv[4] = _fma_pair_5.x; + sv[5] = _fma_pair_5.y; + float2 _fma_pair_6 = fma_f32x2(make_float2(sv[6], sv[7]), _fma_b2_1, _fma_c2_2); + sv[6] = _fma_pair_6.x; + sv[7] = _fma_pair_6.y; + float2 _fma_pair_7 = fma_f32x2(make_float2(sv[8], sv[9]), _fma_b2_1, _fma_c2_2); + sv[8] = _fma_pair_7.x; + sv[9] = _fma_pair_7.y; + float2 _fma_pair_8 = fma_f32x2(make_float2(sv[10], sv[11]), _fma_b2_1, _fma_c2_2); + sv[10] = _fma_pair_8.x; + sv[11] = _fma_pair_8.y; + float2 _fma_pair_9 = fma_f32x2(make_float2(sv[12], sv[13]), _fma_b2_1, _fma_c2_2); + sv[12] = _fma_pair_9.x; + sv[13] = _fma_pair_9.y; + float2 _fma_pair_10 = fma_f32x2(make_float2(sv[14], sv[15]), _fma_b2_1, _fma_c2_2); + sv[14] = _fma_pair_10.x; + sv[15] = _fma_pair_10.y; + float2 _fma_pair_11 = fma_f32x2(make_float2(sv[16], sv[17]), _fma_b2_1, _fma_c2_2); + sv[16] = _fma_pair_11.x; + sv[17] = _fma_pair_11.y; + float2 _fma_pair_12 = fma_f32x2(make_float2(sv[18], sv[19]), _fma_b2_1, _fma_c2_2); + sv[18] = _fma_pair_12.x; + sv[19] = _fma_pair_12.y; + float2 _fma_pair_13 = fma_f32x2(make_float2(sv[20], sv[21]), _fma_b2_1, _fma_c2_2); + sv[20] = _fma_pair_13.x; + sv[21] = _fma_pair_13.y; + float2 _fma_pair_14 = fma_f32x2(make_float2(sv[22], sv[23]), _fma_b2_1, _fma_c2_2); + sv[22] = _fma_pair_14.x; + sv[23] = _fma_pair_14.y; + float2 _fma_pair_15 = fma_f32x2(make_float2(sv[24], sv[25]), _fma_b2_1, _fma_c2_2); + sv[24] = _fma_pair_15.x; + sv[25] = _fma_pair_15.y; + float2 _fma_pair_16 = fma_f32x2(make_float2(sv[26], sv[27]), _fma_b2_1, _fma_c2_2); + sv[26] = _fma_pair_16.x; + sv[27] = _fma_pair_16.y; + float2 _fma_pair_17 = fma_f32x2(make_float2(sv[28], sv[29]), _fma_b2_1, _fma_c2_2); + sv[28] = _fma_pair_17.x; + sv[29] = _fma_pair_17.y; + float2 _fma_pair_18 = fma_f32x2(make_float2(sv[30], sv[31]), _fma_b2_1, _fma_c2_2); + sv[30] = _fma_pair_18.x; + sv[31] = _fma_pair_18.y; + } + } + { + #pragma unroll + for (int _le = 0; _le < 16; _le++) { + if (1 && _le >= 12) { + float2 _exp2_pair_19 = ex2_emulation_f32x2_value(make_float2(sv[_le*2], sv[_le*2 + 1])); + sv[_le*2] = _exp2_pair_19.x; + sv[_le*2 + 1] = _exp2_pair_19.y; + } else { + sv[_le*2] = approx_exp2(sv[_le*2]); + sv[_le*2 + 1] = approx_exp2(sv[_le*2 + 1]); + } + } + } + { + float local_sum = 0.0f; + { + local_sum = sv[0] + sv[4] + sv[8] + sv[12] + sv[16] + sv[20] + sv[24] + sv[28] + (sv[1] + sv[5] + sv[9] + sv[13] + sv[17] + sv[21] + sv[25] + sv[29]) + (sv[2] + sv[6] + sv[10] + sv[14] + sv[18] + sv[22] + sv[26] + sv[30] + (sv[3] + sv[7] + sv[11] + sv[15] + sv[19] + sv[23] + sv[27] + sv[31])); + } + row_sum_val = row_sum_val * acc_scale + local_sum; + } + unsigned int p_pack[8]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv[0]), "f"(sv[1]), + "f"(sv[2]), "f"(sv[3])); + p_pack[0] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv[4]), "f"(sv[5]), + "f"(sv[6]), "f"(sv[7])); + p_pack[1] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv[8]), "f"(sv[9]), + "f"(sv[10]), "f"(sv[11])); + p_pack[2] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv[12]), "f"(sv[13]), + "f"(sv[14]), "f"(sv[15])); + p_pack[3] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv[16]), "f"(sv[17]), + "f"(sv[18]), "f"(sv[19])); + p_pack[4] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv[20]), "f"(sv[21]), + "f"(sv[22]), "f"(sv[23])); + p_pack[5] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv[24]), "f"(sv[25]), + "f"(sv[26]), "f"(sv[27])); + p_pack[6] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv[28]), "f"(sv[29]), + "f"(sv[30]), "f"(sv[31])); + p_pack[7] = _packed; + } + int p_empty_phase = tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + #pragma unroll + for (int px_c = 0; px_c < 2; px_c++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((smem_p_addr + (unsigned int)(phase * 8192) + (unsigned int)(my_row * 128 + (score_half * 64 + px_c * 16) ^ (my_row * 128 + (score_half * 64 + px_c * 16) >> 7 & 7) << 4))), "r"(p_pack[4 * px_c]), "r"(p_pack[4 * px_c + 1]), "r"(p_pack[4 * px_c + 2]), "r"(p_pack[4 * px_c + 3]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + __syncwarp(); + if (elect_sync()) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + } + int final_sum_slot = my_row * 4 + score_half * 2; + smem_stats_sum[final_sum_slot] = row_sum_val; + if (score_half == 0) { + smem_stats_sum[256 + my_row] = row_max_val; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(sum_ready_addr); + const int n_half = warp % 4 / 2; + const int o_tmem_row_base = warp % 4 * 32; + const int corr_row = o_tmem_row_base << 16; + mbarrier_wait(o_done_addr, 0); + unsigned int _phase_sum_ready_0 = 0; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); + _phase_sum_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float fs0 = smem_stats_sum[my_row * 4]; + float fs1 = smem_stats_sum[my_row * 4 + 1]; + float fs2 = smem_stats_sum[my_row * 4 + 2]; + float fs3 = smem_stats_sum[my_row * 4 + 3]; + float fs01 = fs0 + fs1; + float fs23 = fs2 + fs3; + float total_sum = fs01 + fs23; + float safe_sum = ((total_sum > 0.0f) ? total_sum : 1.0f); + float _rcp_0 = approx_rcp(safe_sum); + float output_scale = ((num_split == 1) ? _rcp_0 : 1.0f); + int head_idx = cta_rank * 64 + my_row; + int po_offset = ((batch_idx * num_split + split_idx) * 128 + head_idx) * 512 + n_half * 128; + int direct_offset = (batch_idx * 128 + head_idx) * 512 + n_half * 128; + int epi_base_w = warp * 512; + int epi_wbase = warp * 2048 + lane * 64; + int epi_wkey = lane >> 1 & 3; + int epi_rq = lane >> 2; + int epi_s = lane & 3; + #pragma unroll + for (int vs_local = 1; vs_local < 2; vs_local++) { + int vs = vs_local; + int o_base_epi = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c2 = 0; c2 < 128; c2 += 128) { + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) + : "r"(o_base_epi + c2)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) + : "r"(o_base_epi + c2 + 32)); + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]) + : "r"(o_base_epi + c2 + 64)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]) + : "r"(o_base_epi + c2 + 64 + 32)); + int gmem_base = po_offset + vs * 256 + c2; + { + unsigned int packed_epi[16]; + if (num_split == 1) { + const float2 _scale2_20 = {output_scale * bmm2_scale, output_scale * bmm2_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_0 + 0))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_0 + 0))[_ls], _scale2_20); + } else { + const float2 _scale2_21 = {output_scale, output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_0 + 0))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_0 + 0))[_ls], _scale2_21); + } + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + packed_epi[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int j = 0; j < 16; j += 4) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_epi_addr + (unsigned int)(epi_wbase + (epi_wkey ^ j / 4) * 16)), "r"(packed_epi[j]), "r"(packed_epi[j + 1]), "r"(packed_epi[j + 2]), "r"(packed_epi[j + 3]) : "memory"); + } + __syncwarp(); + #pragma unroll + for (int i = 0; i < 4; i++) { + int epi_r = i * 8 + epi_rq; + unsigned int _smem_epi_reg_0[4]; + { + const unsigned int* _smem_ptr = reinterpret_cast(smem_epi); + #pragma unroll + for (int _lr = 0; _lr < 4; _lr++) + _smem_epi_reg_0[_lr] = _smem_ptr[(epi_base_w + epi_r * 16 + (epi_s ^ epi_r >> 1 & 3) * 4) + _lr]; + } + int epi_head = cta_rank * 64 + tmem_row_base + epi_r; + int epi_col = vs * 256 + c2 + epi_s * 8; + if (num_split == 1) { + reinterpret_cast(O + ((batch_idx * 128 + epi_head) * 512 + n_half * 128 + epi_col))[0] = reinterpret_cast(_smem_epi_reg_0)[0]; + } else { + reinterpret_cast(partial_O + (((batch_idx * num_split + split_idx) * 128 + epi_head) * 512 + n_half * 128 + epi_col))[0] = reinterpret_cast(_smem_epi_reg_0)[0]; + } + } + __syncwarp(); + if (num_split == 1) { + const float2 _scale2_22 = {output_scale * bmm2_scale, output_scale * bmm2_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_0 + 32))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_0 + 32))[_ls], _scale2_22); + } else { + const float2 _scale2_23 = {output_scale, output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_0 + 32))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_0 + 32))[_ls], _scale2_23); + } + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 32], _tmem_load_0[_lp*2+1 + 32])); + packed_epi[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int j_1 = 0; j_1 < 16; j_1 += 4) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_epi_addr + (unsigned int)(epi_wbase + (epi_wkey ^ j_1 / 4) * 16)), "r"(packed_epi[j_1]), "r"(packed_epi[j_1 + 1]), "r"(packed_epi[j_1 + 2]), "r"(packed_epi[j_1 + 3]) : "memory"); + } + __syncwarp(); + #pragma unroll + for (int i_1 = 0; i_1 < 4; i_1++) { + int epi_r_1 = i_1 * 8 + epi_rq; + unsigned int _smem_epi_reg_1[4]; + { + const unsigned int* _smem_ptr = reinterpret_cast(smem_epi); + #pragma unroll + for (int _lr = 0; _lr < 4; _lr++) + _smem_epi_reg_1[_lr] = _smem_ptr[(epi_base_w + epi_r_1 * 16 + (epi_s ^ epi_r_1 >> 1 & 3) * 4) + _lr]; + } + int epi_head_1 = cta_rank * 64 + tmem_row_base + epi_r_1; + int epi_col_1 = vs * 256 + c2 + 32 + epi_s * 8; + if (num_split == 1) { + reinterpret_cast(O + ((batch_idx * 128 + epi_head_1) * 512 + n_half * 128 + epi_col_1))[0] = reinterpret_cast(_smem_epi_reg_1)[0]; + } else { + reinterpret_cast(partial_O + (((batch_idx * num_split + split_idx) * 128 + epi_head_1) * 512 + n_half * 128 + epi_col_1))[0] = reinterpret_cast(_smem_epi_reg_1)[0]; + } + } + __syncwarp(); + } + int gmem_base_0 = po_offset + vs * 256 + c2 + 64; + { + unsigned int packed_epi_1[16]; + if (num_split == 1) { + const float2 _scale2_24 = {output_scale * bmm2_scale, output_scale * bmm2_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_1 + 0))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_1 + 0))[_ls], _scale2_24); + } else { + const float2 _scale2_25 = {output_scale, output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_1 + 0))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_1 + 0))[_ls], _scale2_25); + } + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_1[_lp*2 + 0], _tmem_load_1[_lp*2+1 + 0])); + packed_epi_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int j_2 = 0; j_2 < 16; j_2 += 4) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_epi_addr + (unsigned int)(epi_wbase + (epi_wkey ^ j_2 / 4) * 16)), "r"(packed_epi_1[j_2]), "r"(packed_epi_1[j_2 + 1]), "r"(packed_epi_1[j_2 + 2]), "r"(packed_epi_1[j_2 + 3]) : "memory"); + } + __syncwarp(); + #pragma unroll + for (int i_2 = 0; i_2 < 4; i_2++) { + int epi_r_2 = i_2 * 8 + epi_rq; + unsigned int _smem_epi_reg_2[4]; + { + const unsigned int* _smem_ptr = reinterpret_cast(smem_epi); + #pragma unroll + for (int _lr = 0; _lr < 4; _lr++) + _smem_epi_reg_2[_lr] = _smem_ptr[(epi_base_w + epi_r_2 * 16 + (epi_s ^ epi_r_2 >> 1 & 3) * 4) + _lr]; + } + int epi_head_2 = cta_rank * 64 + tmem_row_base + epi_r_2; + int epi_col_2 = vs * 256 + c2 + 64 + epi_s * 8; + if (num_split == 1) { + reinterpret_cast(O + ((batch_idx * 128 + epi_head_2) * 512 + n_half * 128 + epi_col_2))[0] = reinterpret_cast(_smem_epi_reg_2)[0]; + } else { + reinterpret_cast(partial_O + (((batch_idx * num_split + split_idx) * 128 + epi_head_2) * 512 + n_half * 128 + epi_col_2))[0] = reinterpret_cast(_smem_epi_reg_2)[0]; + } + } + __syncwarp(); + if (num_split == 1) { + const float2 _scale2_26 = {output_scale * bmm2_scale, output_scale * bmm2_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_1 + 32))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_1 + 32))[_ls], _scale2_26); + } else { + const float2 _scale2_27 = {output_scale, output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_1 + 32))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_1 + 32))[_ls], _scale2_27); + } + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_1[_lp*2 + 32], _tmem_load_1[_lp*2+1 + 32])); + packed_epi_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int j_3 = 0; j_3 < 16; j_3 += 4) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_epi_addr + (unsigned int)(epi_wbase + (epi_wkey ^ j_3 / 4) * 16)), "r"(packed_epi_1[j_3]), "r"(packed_epi_1[j_3 + 1]), "r"(packed_epi_1[j_3 + 2]), "r"(packed_epi_1[j_3 + 3]) : "memory"); + } + __syncwarp(); + #pragma unroll + for (int i_3 = 0; i_3 < 4; i_3++) { + int epi_r_3 = i_3 * 8 + epi_rq; + unsigned int _smem_epi_reg_3[4]; + { + const unsigned int* _smem_ptr = reinterpret_cast(smem_epi); + #pragma unroll + for (int _lr = 0; _lr < 4; _lr++) + _smem_epi_reg_3[_lr] = _smem_ptr[(epi_base_w + epi_r_3 * 16 + (epi_s ^ epi_r_3 >> 1 & 3) * 4) + _lr]; + } + int epi_head_3 = cta_rank * 64 + tmem_row_base + epi_r_3; + int epi_col_3 = vs * 256 + c2 + 64 + 32 + epi_s * 8; + if (num_split == 1) { + reinterpret_cast(O + ((batch_idx * 128 + epi_head_3) * 512 + n_half * 128 + epi_col_3))[0] = reinterpret_cast(_smem_epi_reg_3)[0]; + } else { + reinterpret_cast(partial_O + (((batch_idx * num_split + split_idx) * 128 + epi_head_3) * 512 + n_half * 128 + epi_col_3))[0] = reinterpret_cast(_smem_epi_reg_3)[0]; + } + } + __syncwarp(); + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(partial_results_ready_addr); + unsigned int _phase_partial_results_ready_0 = 0; + mbarrier_wait(partial_results_ready_addr, _phase_partial_results_ready_0); + _phase_partial_results_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int o_tmem_row_base_1 = warp % 4 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row_1 = o_tmem_row_base_1 << 16; + float scrub_zero[128]; + #pragma unroll + for (int scrub_i = 0; scrub_i < 128; scrub_i++) { + scrub_zero[scrub_i] = 0.0f; + } + const int scrub_row = (unsigned int)o_tmem_row_base_1 + lane; + #pragma unroll + for (int scrub_col = 0; scrub_col < 512; scrub_col += 128) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x128.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, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %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};" + :: "r"(taddr + (unsigned int)scrub_col + (unsigned int)(scrub_row << 16)), "f"(scrub_zero[0]), "f"(scrub_zero[1]), "f"(scrub_zero[2]), "f"(scrub_zero[3]), "f"(scrub_zero[4]), "f"(scrub_zero[5]), "f"(scrub_zero[6]), "f"(scrub_zero[7]), "f"(scrub_zero[8]), "f"(scrub_zero[9]), "f"(scrub_zero[10]), "f"(scrub_zero[11]), "f"(scrub_zero[12]), "f"(scrub_zero[13]), "f"(scrub_zero[14]), "f"(scrub_zero[15]), "f"(scrub_zero[16]), "f"(scrub_zero[17]), "f"(scrub_zero[18]), "f"(scrub_zero[19]), "f"(scrub_zero[20]), "f"(scrub_zero[21]), "f"(scrub_zero[22]), "f"(scrub_zero[23]), "f"(scrub_zero[24]), "f"(scrub_zero[25]), "f"(scrub_zero[26]), "f"(scrub_zero[27]), "f"(scrub_zero[28]), "f"(scrub_zero[29]), "f"(scrub_zero[30]), "f"(scrub_zero[31]), "f"(scrub_zero[32]), "f"(scrub_zero[33]), "f"(scrub_zero[34]), "f"(scrub_zero[35]), "f"(scrub_zero[36]), "f"(scrub_zero[37]), "f"(scrub_zero[38]), "f"(scrub_zero[39]), "f"(scrub_zero[40]), "f"(scrub_zero[41]), "f"(scrub_zero[42]), "f"(scrub_zero[43]), "f"(scrub_zero[44]), "f"(scrub_zero[45]), "f"(scrub_zero[46]), "f"(scrub_zero[47]), "f"(scrub_zero[48]), "f"(scrub_zero[49]), "f"(scrub_zero[50]), "f"(scrub_zero[51]), "f"(scrub_zero[52]), "f"(scrub_zero[53]), "f"(scrub_zero[54]), "f"(scrub_zero[55]), "f"(scrub_zero[56]), "f"(scrub_zero[57]), "f"(scrub_zero[58]), "f"(scrub_zero[59]), "f"(scrub_zero[60]), "f"(scrub_zero[61]), "f"(scrub_zero[62]), "f"(scrub_zero[63]), "f"(scrub_zero[64]), "f"(scrub_zero[65]), "f"(scrub_zero[66]), "f"(scrub_zero[67]), "f"(scrub_zero[68]), "f"(scrub_zero[69]), "f"(scrub_zero[70]), "f"(scrub_zero[71]), "f"(scrub_zero[72]), "f"(scrub_zero[73]), "f"(scrub_zero[74]), "f"(scrub_zero[75]), "f"(scrub_zero[76]), "f"(scrub_zero[77]), "f"(scrub_zero[78]), "f"(scrub_zero[79]), "f"(scrub_zero[80]), "f"(scrub_zero[81]), "f"(scrub_zero[82]), "f"(scrub_zero[83]), "f"(scrub_zero[84]), "f"(scrub_zero[85]), "f"(scrub_zero[86]), "f"(scrub_zero[87]), "f"(scrub_zero[88]), "f"(scrub_zero[89]), "f"(scrub_zero[90]), "f"(scrub_zero[91]), "f"(scrub_zero[92]), "f"(scrub_zero[93]), "f"(scrub_zero[94]), "f"(scrub_zero[95]), "f"(scrub_zero[96]), "f"(scrub_zero[97]), "f"(scrub_zero[98]), "f"(scrub_zero[99]), "f"(scrub_zero[100]), "f"(scrub_zero[101]), "f"(scrub_zero[102]), "f"(scrub_zero[103]), "f"(scrub_zero[104]), "f"(scrub_zero[105]), "f"(scrub_zero[106]), "f"(scrub_zero[107]), "f"(scrub_zero[108]), "f"(scrub_zero[109]), "f"(scrub_zero[110]), "f"(scrub_zero[111]), "f"(scrub_zero[112]), "f"(scrub_zero[113]), "f"(scrub_zero[114]), "f"(scrub_zero[115]), "f"(scrub_zero[116]), "f"(scrub_zero[117]), "f"(scrub_zero[118]), "f"(scrub_zero[119]), "f"(scrub_zero[120]), "f"(scrub_zero[121]), "f"(scrub_zero[122]), "f"(scrub_zero[123]), "f"(scrub_zero[124]), "f"(scrub_zero[125]), "f"(scrub_zero[126]), "f"(scrub_zero[127])); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_epoch_ready_addr); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_epoch_ready_addr), "r"(cta_rank ^ 1) : "memory"); + int split_idx_1 = blockIdx.x / 2; + int batch_idx_1 = blockIdx.y; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int num_kv_tiles_total_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int blocks_per_split_1 = (num_kv_tiles_total_1 + num_split - 1) / num_split; + int my_start_block_1 = split_idx_1 * blocks_per_split_1; + int my_end_block_raw_1 = my_start_block_1 + blocks_per_split_1; + int my_end_block_1 = ((my_end_block_raw_1 < num_kv_tiles_total_1) ? my_end_block_raw_1 : num_kv_tiles_total_1); + int my_n_tiles_1 = my_end_block_1 - my_start_block_1; + float row_max_val_1 = -MLA_INF; + float row_sum_val_1 = 0.0f; + unsigned int _phase_q_full_0 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < my_n_tiles_1; tile_1++) { + int phase_1 = tile_1 & 1; + int c_wait_phase = tile_1 >> 1 & 1; + float acc_scale_1 = 1.0f; + float sv_1[32]; + mbarrier_wait(s_full_addr + (phase_1) * 8, c_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int s_off_1 = ((phase_1 != 0) ? 128 : 0); + int s_base_1 = taddr + (unsigned int)s_off_1 + 32 + (unsigned int)(tmem_row_base_1 << 16); + int abs_tile_1 = my_start_block_1 + tile_1; + int tail_valid_1 = seqlen_kv_b_1 - abs_tile_1 * 128; + int local_tail_raw_1 = tail_valid_1 - n_half_1 * 64 - 32; + int local_tail_hi_1 = ((local_tail_raw_1 > 32) ? 32 : local_tail_raw_1); + int local_tail_1 = ((local_tail_hi_1 < 0) ? 0 : local_tail_hi_1); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv_1[0]), "=f"(sv_1[1]), "=f"(sv_1[2]), "=f"(sv_1[3]), "=f"(sv_1[4]), "=f"(sv_1[5]), "=f"(sv_1[6]), "=f"(sv_1[7]), "=f"(sv_1[8]), "=f"(sv_1[9]), "=f"(sv_1[10]), "=f"(sv_1[11]), "=f"(sv_1[12]), "=f"(sv_1[13]), "=f"(sv_1[14]), "=f"(sv_1[15]), "=f"(sv_1[16]), "=f"(sv_1[17]), "=f"(sv_1[18]), "=f"(sv_1[19]), "=f"(sv_1[20]), "=f"(sv_1[21]), "=f"(sv_1[22]), "=f"(sv_1[23]), "=f"(sv_1[24]), "=f"(sv_1[25]), "=f"(sv_1[26]), "=f"(sv_1[27]), "=f"(sv_1[28]), "=f"(sv_1[29]), "=f"(sv_1[30]), "=f"(sv_1[31]) + : "r"(s_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + if (local_tail_1 < 32) { + uint32_t _slice_lo_mask_3; + { + int _lim_0 = local_tail_1; + if (_lim_0 <= 0) { _slice_lo_mask_3 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_3 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_3) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_3 & (1u << 0))) sv_1[0] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 1))) sv_1[1] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 2))) sv_1[2] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 3))) sv_1[3] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 4))) sv_1[4] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 5))) sv_1[5] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 6))) sv_1[6] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 7))) sv_1[7] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 8))) sv_1[8] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 9))) sv_1[9] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 10))) sv_1[10] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 11))) sv_1[11] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 12))) sv_1[12] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 13))) sv_1[13] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 14))) sv_1[14] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 15))) sv_1[15] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 16))) sv_1[16] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 17))) sv_1[17] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 18))) sv_1[18] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 19))) sv_1[19] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 20))) sv_1[20] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 21))) sv_1[21] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 22))) sv_1[22] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 23))) sv_1[23] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 24))) sv_1[24] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 25))) sv_1[25] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 26))) sv_1[26] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 27))) sv_1[27] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 28))) sv_1[28] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 29))) sv_1[29] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 30))) sv_1[30] = -MLA_INF; + if (!(_slice_lo_mask_3 & (1u << 31))) sv_1[31] = -MLA_INF; + } + float _max3_15; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_15) : "f"(sv_1[0]), "f"(sv_1[1]), "f"(sv_1[2])); + float _max3_16; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_16) : "f"(sv_1[3]), "f"(sv_1[4]), "f"(sv_1[5])); + float _max3_17; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_17) : "f"(sv_1[6]), "f"(sv_1[7]), "f"(sv_1[8])); + float _max3_18; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_18) : "f"(sv_1[9]), "f"(sv_1[10]), "f"(sv_1[11])); + float _max3_19; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_19) : "f"(sv_1[12]), "f"(sv_1[13]), "f"(sv_1[14])); + float _max3_20; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_20) : "f"(sv_1[15]), "f"(sv_1[16]), "f"(sv_1[17])); + float _max3_21; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_21) : "f"(sv_1[18]), "f"(sv_1[19]), "f"(sv_1[20])); + float _max3_22; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_22) : "f"(sv_1[21]), "f"(sv_1[22]), "f"(sv_1[23])); + float _max3_23; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_23) : "f"(sv_1[24]), "f"(sv_1[25]), "f"(sv_1[26])); + float _max3_24; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_24) : "f"(sv_1[27]), "f"(sv_1[28]), "f"(sv_1[29])); + float _max_7 = max_noftz(sv_1[30], sv_1[31]); + float _max3_25; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_25) : "f"(_max3_15), "f"(_max3_16), "f"(_max3_17)); + float _max3_26; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_26) : "f"(_max3_18), "f"(_max3_19), "f"(_max3_20)); + float _max3_27; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_27) : "f"(_max3_21), "f"(_max3_22), "f"(_max3_23)); + float _max_8 = max_noftz(_max3_24, _max_7); + float _max3_28; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_28) : "f"(_max3_25), "f"(_max3_26), "f"(_max3_27)); + float _max_9 = max_noftz(_max3_28, _max_8); + float local_max_1 = _max_9; + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + __syncwarp(); + if (elect_sync()) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase_1 * 8)), "r"(0) : "memory"); + } + smem_stats_max[phase_1 * 256 + n_half_1 * 128 + 64 + my_row_1] = local_max_1; + if (warp % 2 == 0) { + asm volatile("barrier.sync 6, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 7, 128;" ::: "memory"); + } + asm volatile("tcgen05.fence::after_thread_sync;"); + float mx_s = smem_stats_max[phase_1 * 256 + n_half_1 * 128 + my_row_1]; + float mx_s2_1 = smem_stats_max[phase_1 * 256 + (1 - n_half_1) * 128 + my_row_1]; + float mx_c2_1 = smem_stats_max[phase_1 * 256 + (1 - n_half_1) * 128 + 64 + my_row_1]; + float _max3_29; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + #error "Max3 requires PTX three-input max.f32 support on sm_100+" + #endif + asm volatile("max.f32 %0, %1, %2, %3;" : "=f"(_max3_29) : "f"(local_max_1), "f"(mx_s), "f"(mx_s2_1)); + float new_max_1 = _max3_29; + float _max_10 = max_noftz(new_max_1, mx_c2_1); + new_max_1 = _max_10; + float _max_11 = max_noftz(new_max_1, row_max_val_1); + new_max_1 = _max_11; + float lazy_delta_1 = softmax_scale_log2 * (row_max_val_1 - new_max_1); + new_max_1 = ((lazy_delta_1 < -8.0f) ? new_max_1 : row_max_val_1); + float delta_1 = softmax_scale_log2 * (row_max_val_1 - new_max_1); + float _exp2_2 = approx_exp2(delta_1); + float exp_delta_1 = _exp2_2; + acc_scale_1 = ((row_max_val_1 > -MLA_INF) ? exp_delta_1 : 1.0f); + row_max_val_1 = new_max_1; + int _vote_0 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); + int any_rescale = _vote_0; + if (tile_1 > 0) { + if (any_rescale != 0) { + int o_ready_stage = tile_1 - 1 & 1; + int o_ready_phase = tile_1 - 1 >> 1 & 1; + mbarrier_wait(o_full_addr + (o_ready_stage) * 8, o_ready_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int vs_1 = vs_local_1; + int o_base = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row_1; + #pragma unroll + for (int oc = 0; oc < 2; oc++) { + float _tmem_load_2[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]) + : "r"(o_base + oc * 64)); + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) + : "r"(o_base + oc * 64 + 32)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + { + const float2 _scale2_1 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_1); + const float2 _scale2_2 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_3)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_3)[_ls], _scale2_2); + } + tmem_st_x32_f32(o_base + oc * 64, _tmem_load_2); + tmem_st_x32_f32(o_base + oc * 64 + 32, _tmem_load_3); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + float safe_max_1 = ((new_max_1 == -MLA_INF) ? 0.0f : new_max_1); + float max_scaled_1 = safe_max_1 * softmax_scale_log2; + float neg_max_scaled_1 = -max_scaled_1; + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {neg_max_scaled_1, neg_max_scaled_1}; + float2 _fma_pair_5 = fma_f32x2(make_float2(sv_1[0], sv_1[1]), _fma_b2_3, _fma_c2_4); + sv_1[0] = _fma_pair_5.x; + sv_1[1] = _fma_pair_5.y; + float2 _fma_pair_6 = fma_f32x2(make_float2(sv_1[2], sv_1[3]), _fma_b2_3, _fma_c2_4); + sv_1[2] = _fma_pair_6.x; + sv_1[3] = _fma_pair_6.y; + float2 _fma_pair_7 = fma_f32x2(make_float2(sv_1[4], sv_1[5]), _fma_b2_3, _fma_c2_4); + sv_1[4] = _fma_pair_7.x; + sv_1[5] = _fma_pair_7.y; + float2 _fma_pair_8 = fma_f32x2(make_float2(sv_1[6], sv_1[7]), _fma_b2_3, _fma_c2_4); + sv_1[6] = _fma_pair_8.x; + sv_1[7] = _fma_pair_8.y; + float2 _fma_pair_9 = fma_f32x2(make_float2(sv_1[8], sv_1[9]), _fma_b2_3, _fma_c2_4); + sv_1[8] = _fma_pair_9.x; + sv_1[9] = _fma_pair_9.y; + float2 _fma_pair_10 = fma_f32x2(make_float2(sv_1[10], sv_1[11]), _fma_b2_3, _fma_c2_4); + sv_1[10] = _fma_pair_10.x; + sv_1[11] = _fma_pair_10.y; + float2 _fma_pair_11 = fma_f32x2(make_float2(sv_1[12], sv_1[13]), _fma_b2_3, _fma_c2_4); + sv_1[12] = _fma_pair_11.x; + sv_1[13] = _fma_pair_11.y; + float2 _fma_pair_12 = fma_f32x2(make_float2(sv_1[14], sv_1[15]), _fma_b2_3, _fma_c2_4); + sv_1[14] = _fma_pair_12.x; + sv_1[15] = _fma_pair_12.y; + float2 _fma_pair_13 = fma_f32x2(make_float2(sv_1[16], sv_1[17]), _fma_b2_3, _fma_c2_4); + sv_1[16] = _fma_pair_13.x; + sv_1[17] = _fma_pair_13.y; + float2 _fma_pair_14 = fma_f32x2(make_float2(sv_1[18], sv_1[19]), _fma_b2_3, _fma_c2_4); + sv_1[18] = _fma_pair_14.x; + sv_1[19] = _fma_pair_14.y; + float2 _fma_pair_15 = fma_f32x2(make_float2(sv_1[20], sv_1[21]), _fma_b2_3, _fma_c2_4); + sv_1[20] = _fma_pair_15.x; + sv_1[21] = _fma_pair_15.y; + float2 _fma_pair_16 = fma_f32x2(make_float2(sv_1[22], sv_1[23]), _fma_b2_3, _fma_c2_4); + sv_1[22] = _fma_pair_16.x; + sv_1[23] = _fma_pair_16.y; + float2 _fma_pair_17 = fma_f32x2(make_float2(sv_1[24], sv_1[25]), _fma_b2_3, _fma_c2_4); + sv_1[24] = _fma_pair_17.x; + sv_1[25] = _fma_pair_17.y; + float2 _fma_pair_18 = fma_f32x2(make_float2(sv_1[26], sv_1[27]), _fma_b2_3, _fma_c2_4); + sv_1[26] = _fma_pair_18.x; + sv_1[27] = _fma_pair_18.y; + float2 _fma_pair_19 = fma_f32x2(make_float2(sv_1[28], sv_1[29]), _fma_b2_3, _fma_c2_4); + sv_1[28] = _fma_pair_19.x; + sv_1[29] = _fma_pair_19.y; + float2 _fma_pair_20 = fma_f32x2(make_float2(sv_1[30], sv_1[31]), _fma_b2_3, _fma_c2_4); + sv_1[30] = _fma_pair_20.x; + sv_1[31] = _fma_pair_20.y; + { + #pragma unroll + for (int _le = 0; _le < 16; _le++) { + if (1 && _le >= 12) { + float2 _exp2_pair_21 = ex2_emulation_f32x2_value(make_float2(sv_1[_le*2], sv_1[_le*2 + 1])); + sv_1[_le*2] = _exp2_pair_21.x; + sv_1[_le*2 + 1] = _exp2_pair_21.y; + } else { + sv_1[_le*2] = approx_exp2(sv_1[_le*2]); + sv_1[_le*2 + 1] = approx_exp2(sv_1[_le*2 + 1]); + } + } + } + float local_sum_1 = sv_1[0] + sv_1[4] + sv_1[8] + sv_1[12] + sv_1[16] + sv_1[20] + sv_1[24] + sv_1[28] + (sv_1[1] + sv_1[5] + sv_1[9] + sv_1[13] + sv_1[17] + sv_1[21] + sv_1[25] + sv_1[29]) + (sv_1[2] + sv_1[6] + sv_1[10] + sv_1[14] + sv_1[18] + sv_1[22] + sv_1[26] + sv_1[30] + (sv_1[3] + sv_1[7] + sv_1[11] + sv_1[15] + sv_1[19] + sv_1[23] + sv_1[27] + sv_1[31])); + row_sum_val_1 = row_sum_val_1 * acc_scale_1 + local_sum_1; + unsigned int p_pack_1[8]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_1[0]), "f"(sv_1[1]), + "f"(sv_1[2]), "f"(sv_1[3])); + p_pack_1[0] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_1[4]), "f"(sv_1[5]), + "f"(sv_1[6]), "f"(sv_1[7])); + p_pack_1[1] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_1[8]), "f"(sv_1[9]), + "f"(sv_1[10]), "f"(sv_1[11])); + p_pack_1[2] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_1[12]), "f"(sv_1[13]), + "f"(sv_1[14]), "f"(sv_1[15])); + p_pack_1[3] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_1[16]), "f"(sv_1[17]), + "f"(sv_1[18]), "f"(sv_1[19])); + p_pack_1[4] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_1[20]), "f"(sv_1[21]), + "f"(sv_1[22]), "f"(sv_1[23])); + p_pack_1[5] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_1[24]), "f"(sv_1[25]), + "f"(sv_1[26]), "f"(sv_1[27])); + p_pack_1[6] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_1[28]), "f"(sv_1[29]), + "f"(sv_1[30]), "f"(sv_1[31])); + p_pack_1[7] = _packed; + } + int p_empty_phase_1 = tile_1 >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase_1) * 8, p_empty_phase_1); + #pragma unroll + for (int px_c_1 = 0; px_c_1 < 2; px_c_1++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((smem_p_addr + (unsigned int)(phase_1 * 8192) + (unsigned int)(my_row_1 * 128 + (n_half_1 * 64 + 32 + px_c_1 * 16) ^ (my_row_1 * 128 + (n_half_1 * 64 + 32 + px_c_1 * 16) >> 7 & 7) << 4))), "r"(p_pack_1[4 * px_c_1]), "r"(p_pack_1[4 * px_c_1 + 1]), "r"(p_pack_1[4 * px_c_1 + 2]), "r"(p_pack_1[4 * px_c_1 + 3]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncwarp(); + if (elect_sync()) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase_1 * 8)), "r"(0) : "memory"); + } + } + { + mbarrier_wait(o_done_addr, 0); + } + smem_stats_sum[my_row_1 * 4 + n_half_1 * 2 + 1] = row_sum_val_1; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(sum_ready_addr); + unsigned int _phase_sum_ready_0_1 = 0; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0_1); + _phase_sum_ready_0_1 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float fs0_1 = smem_stats_sum[my_row_1 * 4]; + float fs1_1 = smem_stats_sum[my_row_1 * 4 + 1]; + float fs2_1 = smem_stats_sum[my_row_1 * 4 + 2]; + float fs3_1 = smem_stats_sum[my_row_1 * 4 + 3]; + float fs01_1 = fs0_1 + fs1_1; + float fs23_1 = fs2_1 + fs3_1; + float total_sum_1 = fs01_1 + fs23_1; + float final_row_max = smem_stats_sum[256 + my_row_1]; + float safe_sum_1 = ((total_sum_1 > 0.0f) ? total_sum_1 : 1.0f); + float _rcp_1 = approx_rcp(safe_sum_1); + float output_scale_1 = ((num_split == 1) ? _rcp_1 : 1.0f); + int head_idx_1 = cta_rank * 64 + my_row_1; + if (n_half_1 == 0) { + int stat_off = ((batch_idx_1 * num_split + split_idx_1) * 128 + head_idx_1) * 2; + float row_max_scaled = final_row_max * softmax_scale_log2; + float stored_max = ((final_row_max == -MLA_INF) ? -MLA_INF : row_max_scaled); + float packed_stats[2]; + packed_stats[0] = stored_max; + packed_stats[1] = total_sum_1; + { + float2 _v2 = make_float2(packed_stats[0 + 0], packed_stats[0 + 1]); + *reinterpret_cast(partial_stats + stat_off + 0) = _v2; + } + if (num_split == 1) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(total_sum_1)); + float direct_lse = ((total_sum_1 > 0.0f) ? stored_max + _log2_0 : -MLA_INF); + LSE[batch_idx_1 * 128 + head_idx_1] = direct_lse; + } + } + int po_offset_1 = ((batch_idx_1 * num_split + split_idx_1) * 128 + head_idx_1) * 512 + n_half_1 * 128; + int direct_offset_1 = (batch_idx_1 * 128 + head_idx_1) * 512 + n_half_1 * 128; + int epi_base_w_1 = warp * 512; + int epi_wbase_1 = warp * 2048 + lane * 64; + int epi_wkey_1 = lane >> 1 & 3; + int epi_rq_1 = lane >> 2; + int epi_s_1 = lane & 3; + #pragma unroll + for (int vs_local_2 = 0; vs_local_2 < 1; vs_local_2++) { + int vs_2 = vs_local_2; + int o_base_epi_1 = taddr + 256 + (unsigned int)(vs_2 * 128) + (unsigned int)corr_row_1; + #pragma unroll + for (int c2_1 = 0; c2_1 < 128; c2_1 += 128) { + float _tmem_load_4[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]), "=f"(_tmem_load_4[2]), "=f"(_tmem_load_4[3]), "=f"(_tmem_load_4[4]), "=f"(_tmem_load_4[5]), "=f"(_tmem_load_4[6]), "=f"(_tmem_load_4[7]), "=f"(_tmem_load_4[8]), "=f"(_tmem_load_4[9]), "=f"(_tmem_load_4[10]), "=f"(_tmem_load_4[11]), "=f"(_tmem_load_4[12]), "=f"(_tmem_load_4[13]), "=f"(_tmem_load_4[14]), "=f"(_tmem_load_4[15]), "=f"(_tmem_load_4[16]), "=f"(_tmem_load_4[17]), "=f"(_tmem_load_4[18]), "=f"(_tmem_load_4[19]), "=f"(_tmem_load_4[20]), "=f"(_tmem_load_4[21]), "=f"(_tmem_load_4[22]), "=f"(_tmem_load_4[23]), "=f"(_tmem_load_4[24]), "=f"(_tmem_load_4[25]), "=f"(_tmem_load_4[26]), "=f"(_tmem_load_4[27]), "=f"(_tmem_load_4[28]), "=f"(_tmem_load_4[29]), "=f"(_tmem_load_4[30]), "=f"(_tmem_load_4[31]) + : "r"(o_base_epi_1 + c2_1)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_4[32]), "=f"(_tmem_load_4[33]), "=f"(_tmem_load_4[34]), "=f"(_tmem_load_4[35]), "=f"(_tmem_load_4[36]), "=f"(_tmem_load_4[37]), "=f"(_tmem_load_4[38]), "=f"(_tmem_load_4[39]), "=f"(_tmem_load_4[40]), "=f"(_tmem_load_4[41]), "=f"(_tmem_load_4[42]), "=f"(_tmem_load_4[43]), "=f"(_tmem_load_4[44]), "=f"(_tmem_load_4[45]), "=f"(_tmem_load_4[46]), "=f"(_tmem_load_4[47]), "=f"(_tmem_load_4[48]), "=f"(_tmem_load_4[49]), "=f"(_tmem_load_4[50]), "=f"(_tmem_load_4[51]), "=f"(_tmem_load_4[52]), "=f"(_tmem_load_4[53]), "=f"(_tmem_load_4[54]), "=f"(_tmem_load_4[55]), "=f"(_tmem_load_4[56]), "=f"(_tmem_load_4[57]), "=f"(_tmem_load_4[58]), "=f"(_tmem_load_4[59]), "=f"(_tmem_load_4[60]), "=f"(_tmem_load_4[61]), "=f"(_tmem_load_4[62]), "=f"(_tmem_load_4[63]) + : "r"(o_base_epi_1 + c2_1 + 32)); + float _tmem_load_5[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]) + : "r"(o_base_epi_1 + c2_1 + 64)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_5[32]), "=f"(_tmem_load_5[33]), "=f"(_tmem_load_5[34]), "=f"(_tmem_load_5[35]), "=f"(_tmem_load_5[36]), "=f"(_tmem_load_5[37]), "=f"(_tmem_load_5[38]), "=f"(_tmem_load_5[39]), "=f"(_tmem_load_5[40]), "=f"(_tmem_load_5[41]), "=f"(_tmem_load_5[42]), "=f"(_tmem_load_5[43]), "=f"(_tmem_load_5[44]), "=f"(_tmem_load_5[45]), "=f"(_tmem_load_5[46]), "=f"(_tmem_load_5[47]), "=f"(_tmem_load_5[48]), "=f"(_tmem_load_5[49]), "=f"(_tmem_load_5[50]), "=f"(_tmem_load_5[51]), "=f"(_tmem_load_5[52]), "=f"(_tmem_load_5[53]), "=f"(_tmem_load_5[54]), "=f"(_tmem_load_5[55]), "=f"(_tmem_load_5[56]), "=f"(_tmem_load_5[57]), "=f"(_tmem_load_5[58]), "=f"(_tmem_load_5[59]), "=f"(_tmem_load_5[60]), "=f"(_tmem_load_5[61]), "=f"(_tmem_load_5[62]), "=f"(_tmem_load_5[63]) + : "r"(o_base_epi_1 + c2_1 + 64 + 32)); + int gmem_base_1 = po_offset_1 + vs_2 * 256 + c2_1; + { + unsigned int packed_epi_2[16]; + if (num_split == 1) { + const float2 _scale2_22 = {output_scale_1 * bmm2_scale, output_scale_1 * bmm2_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_4 + 0))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_4 + 0))[_ls], _scale2_22); + } else { + const float2 _scale2_23 = {output_scale_1, output_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_4 + 0))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_4 + 0))[_ls], _scale2_23); + } + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 0], _tmem_load_4[_lp*2+1 + 0])); + packed_epi_2[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int j_4 = 0; j_4 < 16; j_4 += 4) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_epi_addr + (unsigned int)(epi_wbase_1 + (epi_wkey_1 ^ j_4 / 4) * 16)), "r"(packed_epi_2[j_4]), "r"(packed_epi_2[j_4 + 1]), "r"(packed_epi_2[j_4 + 2]), "r"(packed_epi_2[j_4 + 3]) : "memory"); + } + __syncwarp(); + #pragma unroll + for (int i_4 = 0; i_4 < 4; i_4++) { + int epi_r_4 = i_4 * 8 + epi_rq_1; + unsigned int _smem_epi_reg_4[4]; + { + const unsigned int* _smem_ptr = reinterpret_cast(smem_epi); + #pragma unroll + for (int _lr = 0; _lr < 4; _lr++) + _smem_epi_reg_4[_lr] = _smem_ptr[(epi_base_w_1 + epi_r_4 * 16 + (epi_s_1 ^ epi_r_4 >> 1 & 3) * 4) + _lr]; + } + int epi_head_4 = cta_rank * 64 + tmem_row_base_1 + epi_r_4; + int epi_col_4 = vs_2 * 256 + c2_1 + epi_s_1 * 8; + if (num_split == 1) { + reinterpret_cast(O + ((batch_idx_1 * 128 + epi_head_4) * 512 + n_half_1 * 128 + epi_col_4))[0] = reinterpret_cast(_smem_epi_reg_4)[0]; + } else { + reinterpret_cast(partial_O + (((batch_idx_1 * num_split + split_idx_1) * 128 + epi_head_4) * 512 + n_half_1 * 128 + epi_col_4))[0] = reinterpret_cast(_smem_epi_reg_4)[0]; + } + } + __syncwarp(); + if (num_split == 1) { + const float2 _scale2_24 = {output_scale_1 * bmm2_scale, output_scale_1 * bmm2_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_4 + 32))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_4 + 32))[_ls], _scale2_24); + } else { + const float2 _scale2_25 = {output_scale_1, output_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_4 + 32))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_4 + 32))[_ls], _scale2_25); + } + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 32], _tmem_load_4[_lp*2+1 + 32])); + packed_epi_2[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int j_5 = 0; j_5 < 16; j_5 += 4) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_epi_addr + (unsigned int)(epi_wbase_1 + (epi_wkey_1 ^ j_5 / 4) * 16)), "r"(packed_epi_2[j_5]), "r"(packed_epi_2[j_5 + 1]), "r"(packed_epi_2[j_5 + 2]), "r"(packed_epi_2[j_5 + 3]) : "memory"); + } + __syncwarp(); + #pragma unroll + for (int i_5 = 0; i_5 < 4; i_5++) { + int epi_r_5 = i_5 * 8 + epi_rq_1; + unsigned int _smem_epi_reg_5[4]; + { + const unsigned int* _smem_ptr = reinterpret_cast(smem_epi); + #pragma unroll + for (int _lr = 0; _lr < 4; _lr++) + _smem_epi_reg_5[_lr] = _smem_ptr[(epi_base_w_1 + epi_r_5 * 16 + (epi_s_1 ^ epi_r_5 >> 1 & 3) * 4) + _lr]; + } + int epi_head_5 = cta_rank * 64 + tmem_row_base_1 + epi_r_5; + int epi_col_5 = vs_2 * 256 + c2_1 + 32 + epi_s_1 * 8; + if (num_split == 1) { + reinterpret_cast(O + ((batch_idx_1 * 128 + epi_head_5) * 512 + n_half_1 * 128 + epi_col_5))[0] = reinterpret_cast(_smem_epi_reg_5)[0]; + } else { + reinterpret_cast(partial_O + (((batch_idx_1 * num_split + split_idx_1) * 128 + epi_head_5) * 512 + n_half_1 * 128 + epi_col_5))[0] = reinterpret_cast(_smem_epi_reg_5)[0]; + } + } + __syncwarp(); + } + int gmem_base_0_1 = po_offset_1 + vs_2 * 256 + c2_1 + 64; + { + unsigned int packed_epi_3[16]; + if (num_split == 1) { + const float2 _scale2_26 = {output_scale_1 * bmm2_scale, output_scale_1 * bmm2_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_5 + 0))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_5 + 0))[_ls], _scale2_26); + } else { + const float2 _scale2_27 = {output_scale_1, output_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_5 + 0))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_5 + 0))[_ls], _scale2_27); + } + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_5[_lp*2 + 0], _tmem_load_5[_lp*2+1 + 0])); + packed_epi_3[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int j_6 = 0; j_6 < 16; j_6 += 4) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_epi_addr + (unsigned int)(epi_wbase_1 + (epi_wkey_1 ^ j_6 / 4) * 16)), "r"(packed_epi_3[j_6]), "r"(packed_epi_3[j_6 + 1]), "r"(packed_epi_3[j_6 + 2]), "r"(packed_epi_3[j_6 + 3]) : "memory"); + } + __syncwarp(); + #pragma unroll + for (int i_6 = 0; i_6 < 4; i_6++) { + int epi_r_6 = i_6 * 8 + epi_rq_1; + unsigned int _smem_epi_reg_6[4]; + { + const unsigned int* _smem_ptr = reinterpret_cast(smem_epi); + #pragma unroll + for (int _lr = 0; _lr < 4; _lr++) + _smem_epi_reg_6[_lr] = _smem_ptr[(epi_base_w_1 + epi_r_6 * 16 + (epi_s_1 ^ epi_r_6 >> 1 & 3) * 4) + _lr]; + } + int epi_head_6 = cta_rank * 64 + tmem_row_base_1 + epi_r_6; + int epi_col_6 = vs_2 * 256 + c2_1 + 64 + epi_s_1 * 8; + if (num_split == 1) { + reinterpret_cast(O + ((batch_idx_1 * 128 + epi_head_6) * 512 + n_half_1 * 128 + epi_col_6))[0] = reinterpret_cast(_smem_epi_reg_6)[0]; + } else { + reinterpret_cast(partial_O + (((batch_idx_1 * num_split + split_idx_1) * 128 + epi_head_6) * 512 + n_half_1 * 128 + epi_col_6))[0] = reinterpret_cast(_smem_epi_reg_6)[0]; + } + } + __syncwarp(); + if (num_split == 1) { + const float2 _scale2_28 = {output_scale_1 * bmm2_scale, output_scale_1 * bmm2_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_5 + 32))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_5 + 32))[_ls], _scale2_28); + } else { + const float2 _scale2_29 = {output_scale_1, output_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_5 + 32))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_5 + 32))[_ls], _scale2_29); + } + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_5[_lp*2 + 32], _tmem_load_5[_lp*2+1 + 32])); + packed_epi_3[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int j_7 = 0; j_7 < 16; j_7 += 4) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_epi_addr + (unsigned int)(epi_wbase_1 + (epi_wkey_1 ^ j_7 / 4) * 16)), "r"(packed_epi_3[j_7]), "r"(packed_epi_3[j_7 + 1]), "r"(packed_epi_3[j_7 + 2]), "r"(packed_epi_3[j_7 + 3]) : "memory"); + } + __syncwarp(); + #pragma unroll + for (int i_7 = 0; i_7 < 4; i_7++) { + int epi_r_7 = i_7 * 8 + epi_rq_1; + unsigned int _smem_epi_reg_7[4]; + { + const unsigned int* _smem_ptr = reinterpret_cast(smem_epi); + #pragma unroll + for (int _lr = 0; _lr < 4; _lr++) + _smem_epi_reg_7[_lr] = _smem_ptr[(epi_base_w_1 + epi_r_7 * 16 + (epi_s_1 ^ epi_r_7 >> 1 & 3) * 4) + _lr]; + } + int epi_head_7 = cta_rank * 64 + tmem_row_base_1 + epi_r_7; + int epi_col_7 = vs_2 * 256 + c2_1 + 64 + 32 + epi_s_1 * 8; + if (num_split == 1) { + reinterpret_cast(O + ((batch_idx_1 * 128 + epi_head_7) * 512 + n_half_1 * 128 + epi_col_7))[0] = reinterpret_cast(_smem_epi_reg_7)[0]; + } else { + reinterpret_cast(partial_O + (((batch_idx_1 * num_split + split_idx_1) * 128 + epi_head_7) * 512 + n_half_1 * 128 + epi_col_7))[0] = reinterpret_cast(_smem_epi_reg_7)[0]; + } + } + __syncwarp(); + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + __syncwarp(); + if (elect_sync()) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + __syncwarp(); + if (elect_sync()) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr + 8), "r"(0) : "memory"); + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(partial_results_ready_addr); + unsigned int _phase_partial_results_ready_0_1 = 0; + mbarrier_wait(partial_results_ready_addr, _phase_partial_results_ready_0_1); + _phase_partial_results_ready_0_1 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + if (USE_PDL && num_split > 1) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + asm volatile("barrier.sync 13, 160;" ::: "memory"); + if (warp == 4) { + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int wg3_dummy = 0; + unsigned int mma_k_stage = 0; + unsigned int mma_v_stage = 0; + int split_idx_2 = blockIdx.x / 2; + int batch_idx_2 = blockIdx.y; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int num_kv_tiles_total_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int blocks_per_split_2 = (num_kv_tiles_total_2 + num_split - 1) / num_split; + int my_start_block_2 = split_idx_2 * blocks_per_split_2; + int my_end_block_raw_2 = my_start_block_2 + blocks_per_split_2; + int my_end_block_2 = ((my_end_block_raw_2 < num_kv_tiles_total_2) ? my_end_block_raw_2 : num_kv_tiles_total_2); + int my_n_tiles_2 = my_end_block_2 - my_start_block_2; + unsigned int _phase_tmem_epoch_ready_0 = 0; + mbarrier_wait(tmem_epoch_ready_addr, _phase_tmem_epoch_ready_0); + _phase_tmem_epoch_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0_1 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0_1); + _phase_q_full_0_1 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(q_pair_ready_addr), "r"(0) : "memory"); + unsigned int _phase_q_pair_ready_0 = 0; + if (cta_rank == 0) { + mbarrier_wait_cluster_hint(q_pair_ready_addr, _phase_q_pair_ready_0, 10000000); + _phase_q_pair_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + int first_pv = 1; + unsigned int _phase_k_full = 0; + unsigned int _phase_v_full = 0; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < my_n_tiles_2 + 1; tile_2++) { + int phase_2 = tile_2 & 1; + int s_empty_phase = tile_2 >> 1 & 1 ^ 1; + if (my_n_tiles_2 > tile_2) { + int score_col = ((phase_2 != 0) ? 128 : 0); + if (cta_rank == 0) { + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + { + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (0) * 2560; + int _mma_b_lo_0 = (((smem_k_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 2560; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(0)); + } + { + int _mma_a_lo_2 = (((smem_q_addr + 8192) >> 4) & 0x3FFF) + (0) * 2560; + int _mma_b_lo_2 = (((smem_k_addr + 8192) >> 4) & 0x3FFF) + (mma_k_stage) * 2560; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + { + int _mma_a_lo_4 = (((smem_q_addr + 16384) >> 4) & 0x3FFF) + (0) * 2560; + int _mma_b_lo_4 = (((smem_k_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 2560; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + { + int _mma_a_lo_6 = (((smem_q_addr + 24576) >> 4) & 0x3FFF) + (0) * 2560; + int _mma_b_lo_6 = (((smem_k_addr + 24576) >> 4) & 0x3FFF) + (mma_k_stage) * 2560; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + { + int _mma_a_lo_9 = (((smem_q_addr + 32768) >> 4) & 0x3FFF) + (0) * 2560; + int _mma_b_lo_9 = (((smem_k_addr + 32768) >> 4) & 0x3FFF) + (mma_k_stage) * 2560; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_9), "r"(_mma_b_lo_9), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(k_empty_addr + (mma_k_stage) * 8, (uint16_t)(3)); + } + mma_k_stage += 1; + if (mma_k_stage == 2) { mma_k_stage = 0; _phase_k_full ^= 1; } + } + if (tile_2 > 1) { + int prev_phase = tile_2 - 2 & 1; + int prev_tile = tile_2 - 2; + int pv_wait_phase = prev_tile >> 1 & 1; + if (cta_rank == 0) { + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_10 = (((smem_p_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; + int _mma_b_lo_10 = ((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138477584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_11 = (((smem_p_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; + int _mma_b_lo_11 = ((((smem_v_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138477584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); + elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); + mma_v_stage += 1; + if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr + (prev_phase) * 8, (uint16_t)(3)); + } + first_pv = 0; + } + } + int last_phase = my_n_tiles_2 - 1 & 1; + int last_tile = my_n_tiles_2 - 1; + int drain_wait_phase = last_tile >> 1 & 1; + if (cta_rank == 0) { + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_12 = (((smem_p_addr) >> 4) & 0x3FFF) + (last_phase) * 512; + int _mma_b_lo_12 = ((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138477584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_12), "r"(_mma_b_lo_12), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_13 = (((smem_p_addr) >> 4) & 0x3FFF) + (last_phase) * 512; + int _mma_b_lo_13 = ((((smem_v_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138477584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_13), "r"(_mma_b_lo_13), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); + elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); + mma_v_stage += 1; + if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_done_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mbarrier_wait(o_empty_addr, 0); + mbarrier_wait(o_empty_addr + 8, 0); + } + asm volatile("barrier.arrive 13, 160;" ::: "memory"); + } + } + // ---- Role: page_offsets_warp ---- + if (warp == 9) { + { // page_offsets_warp_main + const int wg3_dummy_1 = 0; + unsigned int _phase_q_full_0_2 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0_2); + _phase_q_full_0_2 ^= 1; + } + } + // ---- Role: padding_warp ---- + if (warp == 10) { + { // padding_warp_main + const int wg3_dummy_2 = 0; + unsigned int _phase_q_full_0_3 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0_3); + _phase_q_full_0_3 ^= 1; + } + } + // ---- Role: load_warp ---- + if (warp == 11) { + { // load_warp_main + const int wg3_dummy_3 = 0; + unsigned int load_k_stage = 0; + unsigned int load_v_stage = 0; + int split_idx_3 = blockIdx.x / 2; + int batch_idx_3 = blockIdx.y; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int num_kv_tiles_total_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int blocks_per_split_3 = (num_kv_tiles_total_3 + num_split - 1) / num_split; + int my_start_block_3 = split_idx_3 * blocks_per_split_3; + int my_end_block_raw_3 = my_start_block_3 + blocks_per_split_3; + int my_end_block_3 = ((my_end_block_raw_3 < num_kv_tiles_total_3) ? my_end_block_raw_3 : num_kv_tiles_total_3); + int my_n_tiles_3 = my_end_block_3 - my_start_block_3; + int page_batch_idx = work_batch_indices[batch_idx_3]; + int pt_base = page_batch_idx * max_pages_per_seq; + int q_row_global = batch_idx_3 * 128 + cta_rank * 64; + unsigned int _phase_q_empty_0 = 1; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 40960); + #pragma unroll + for (int s = 0; s < 5; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_2d_gmem2smem(q_dst, tmap_q, s * 128, q_row_global, q_full_addr); + } + } + int pg_k0 = page_table[pt_base + 2 * my_start_block_3 + cta_rank]; + unsigned int _phase_k_empty = 1; + mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(81920)) : "memory"); + } + } + if (elect_sync()) { + int k0_dst = smem_k_addr + load_k_stage * 40960; + #pragma unroll + for (int n = 0; n < 5; n++) { + tma_2d_gmem2smem_cta2(k0_dst + n * 8192, tmap_k, n * 128, pg_k0 * 64, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF)); + } + } + load_k_stage += 1; + if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } + unsigned int _phase_v_empty = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < my_n_tiles_3; tile_3++) { + if (my_n_tiles_3 > tile_3 + 1) { + int abs_tile_k = my_start_block_3 + tile_3 + 1; + int pg_k = page_table[pt_base + 2 * abs_tile_k + cta_rank]; + mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(81920)) : "memory"); + } + } + if (elect_sync()) { + int dst_k = smem_k_addr + load_k_stage * 40960; + #pragma unroll + for (int n_1 = 0; n_1 < 5; n_1++) { + tma_2d_gmem2smem_cta2(dst_k + n_1 * 8192, tmap_k, n_1 * 128, pg_k * 64, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF)); + } + } + load_k_stage += 1; + if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } + } + if (tile_3 > 0) { + int abs_tile_v = my_start_block_3 + (tile_3 - 1); + int pg_v0 = page_table[pt_base + 2 * abs_tile_v]; + int pg_v1 = page_table[pt_base + 2 * abs_tile_v + 1]; + mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); + } + } + if (elect_sync()) { + int v_dst = smem_v_addr + load_v_stage * 32768; + int v0_z = cta_rank; + int v1_z = 2 + cta_rank; + tma_2d_gmem2smem_cta2(v_dst, tmap_v, v0_z * 128, pg_v0 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + tma_2d_gmem2smem_cta2(v_dst + 16384, tmap_v, v1_z * 128, pg_v0 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + tma_2d_gmem2smem_cta2(v_dst + 8192, tmap_v, v0_z * 128, pg_v1 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + tma_2d_gmem2smem_cta2(v_dst + 16384 + 8192, tmap_v, v1_z * 128, pg_v1 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + } + load_v_stage += 1; + if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } + } + } + int abs_tile_last = my_start_block_3 + (my_n_tiles_3 - 1); + int pg_vl0 = page_table[pt_base + 2 * abs_tile_last]; + int pg_vl1 = page_table[pt_base + 2 * abs_tile_last + 1]; + mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); + } + } + if (elect_sync()) { + int v_dst_last = smem_v_addr + load_v_stage * 32768; + int v0_z_last = cta_rank; + int v1_z_last = 2 + cta_rank; + tma_2d_gmem2smem_cta2(v_dst_last, tmap_v, v0_z_last * 128, pg_vl0 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + tma_2d_gmem2smem_cta2(v_dst_last + 16384, tmap_v, v1_z_last * 128, pg_vl0 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + tma_2d_gmem2smem_cta2(v_dst_last + 8192, tmap_v, v0_z_last * 128, pg_vl1 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + tma_2d_gmem2smem_cta2(v_dst_last + 16384 + 8192, tmap_v, v1_z_last * 128, pg_vl1 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + } + load_v_stage += 1; + if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_fp8_p32_exact_producer_105fd84009.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_fp8_p32_exact_producer_105fd84009.cu new file mode 100644 index 00000000000..dbdce10dc8f --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_fp8_p32_exact_producer_105fd84009.cu @@ -0,0 +1,3201 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_OFFSET 0 +#define NUM_Q_PIPE_STAGES 2 +#define NUM_KV_PIPE_STAGES 9 +#define NUM_PAGE_PIPE_STAGES 6 +#define NUM_ORDER_PIPE_STAGES 1 +#define SMEM_SMEM_Q0_OFF 0 +#define SMEM_SMEM_Q0_STAGE_BYTES 2048 +#define SMEM_SMEM_Q0_STRIDE 10240 +#define SMEM_SMEM_Q1_OFF 2048 +#define SMEM_SMEM_Q1_STAGE_BYTES 2048 +#define SMEM_SMEM_Q1_STRIDE 10240 +#define SMEM_SMEM_Q2_OFF 4096 +#define SMEM_SMEM_Q2_STAGE_BYTES 2048 +#define SMEM_SMEM_Q2_STRIDE 10240 +#define SMEM_SMEM_Q3_OFF 6144 +#define SMEM_SMEM_Q3_STAGE_BYTES 2048 +#define SMEM_SMEM_Q3_STRIDE 10240 +#define SMEM_SMEM_Q4_OFF 8192 +#define SMEM_SMEM_Q4_STAGE_BYTES 2048 +#define SMEM_SMEM_Q4_STRIDE 10240 +#define SMEM_SMEM_KV_OFF 20480 +#define SMEM_SMEM_KV_STAGE_BYTES 16384 +#define SMEM_SMEM_KV_STRIDE 16384 +#define SMEM_SMEM_V_OFF 20480 +#define SMEM_SMEM_V_STAGE_BYTES 16384 +#define SMEM_SMEM_V_STRIDE 16384 +#define SMEM_SMEM_P0_OFF 167936 +#define SMEM_SMEM_P0_STAGE_BYTES 2048 +#define SMEM_SMEM_P0_STRIDE 2048 +#define SMEM_SMEM_P1_OFF 169984 +#define SMEM_SMEM_P1_STAGE_BYTES 2048 +#define SMEM_SMEM_P1_STRIDE 2048 +#define SMEM_PAGE_OFFSETS_OFF 172032 +#define SMEM_PAGE_OFFSETS_STAGE_BYTES 128 +#define SMEM_PAGE_OFFSETS_STRIDE 128 +#define SMEM_EXCHANGE0_OFF 172800 +#define SMEM_EXCHANGE0_STAGE_BYTES 256 +#define SMEM_EXCHANGE0_STRIDE 256 +#define SMEM_EXCHANGE1_OFF 173056 +#define SMEM_EXCHANGE1_STAGE_BYTES 256 +#define SMEM_EXCHANGE1_STRIDE 256 +#define SMEM_EXCHANGE0_U32_OFF 172800 +#define SMEM_EXCHANGE0_U32_STAGE_BYTES 256 +#define SMEM_EXCHANGE0_U32_STRIDE 256 +#define SMEM_EXCHANGE1_U32_OFF 173056 +#define SMEM_EXCHANGE1_U32_STAGE_BYTES 256 +#define SMEM_EXCHANGE1_U32_STRIDE 256 +#define SMEM_CORR_STATS_REDUCE_OFF 173312 +#define SMEM_CORR_STATS_REDUCE_STAGE_BYTES 256 +#define SMEM_CORR_STATS_REDUCE_STRIDE 256 +#define SMEM_REDUCTION_TICKET_OFF 173696 +#define SMEM_REDUCTION_TICKET_STAGE_BYTES 4 +#define SMEM_REDUCTION_TICKET_STRIDE 4 +#define SMEM_SMEM_O_OFF 173824 +#define SMEM_SMEM_O_STAGE_BYTES 4096 +#define SMEM_SMEM_O_STRIDE 4096 +#define SMEM_TOTAL 178560 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7])); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void elect_commit2(int mbar_addr0, int mbar_addr1) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%1];\n\t" + "}\n" + :: "r"(mbar_addr0), "r"(mbar_addr1) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit2(int mbar_addr0, int mbar_addr1) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%1];\n\t" + :: "r"(mbar_addr0), "r"(mbar_addr1) : "memory"); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x8(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x8.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x8_wait(float* dst, int addr) { + tmem_ld_x8(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_fp8_p32_exact_producer(const __grid_constant__ CUtensorMap Q_map, const __grid_constant__ CUtensorMap K_map, const __grid_constant__ CUtensorMap V_map, int* __restrict__ page_table, int* __restrict__ seq_lens, __nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_stats, __nv_bfloat16* __restrict__ O, unsigned int* __restrict__ completion, int q_len, int page_table_stride, int max_num_ctas_q, int max_num_ctas_kv, float bmm1_scale_log2, float bmm2_scale) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem + 177920; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 16) + #define kv_full_addr (mbar_base + 32) + #define kv_empty_addr (mbar_base + 104) + #define s_full0_addr (mbar_base + 176) + #define s_full1_addr (mbar_base + 184) + #define _source_p_empty0_addr (mbar_base + 192) + #define _source_p_empty1_addr (mbar_base + 200) + #define o_full0_addr (mbar_base + 208) + #define o_full1_addr (mbar_base + 216) + #define page_full_addr (mbar_base + 224) + #define page_empty_addr (mbar_base + 272) + #define s_empty0_addr (mbar_base + 320) + #define scale_ready0_addr (mbar_base + 328) + #define stats_empty0_addr (mbar_base + 336) + #define s_empty1_addr (mbar_base + 344) + #define scale_ready1_addr (mbar_base + 352) + #define stats_empty1_addr (mbar_base + 360) + #define order_p01_0_addr (mbar_base + 368) + #define order_p01_1_addr (mbar_base + 376) + #define _source_p_ready0_addr (mbar_base + 384) + #define _source_p_ready1_addr (mbar_base + 392) + #define o_empty0_addr (mbar_base + 400) + #define o_empty1_addr (mbar_base + 408) + #define _source_barrier_reservation_addr (mbar_base + 416) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + uint8_t* smem_q0 = reinterpret_cast(smem_raw + 0); + const int smem_q0_addr = smem + 0; + uint8_t* smem_q1 = reinterpret_cast(smem_raw + 2048); + const int smem_q1_addr = smem + 2048; + uint8_t* smem_q2 = reinterpret_cast(smem_raw + 4096); + const int smem_q2_addr = smem + 4096; + uint8_t* smem_q3 = reinterpret_cast(smem_raw + 6144); + const int smem_q3_addr = smem + 6144; + uint8_t* smem_q4 = reinterpret_cast(smem_raw + 8192); + const int smem_q4_addr = smem + 8192; + uint8_t* smem_kv = reinterpret_cast(smem_raw + 20480); + const int smem_kv_addr = smem + 20480; + uint8_t* smem_v = reinterpret_cast(smem_raw + 20480); + const int smem_v_addr = smem + 20480; + uint8_t* smem_p0 = reinterpret_cast(smem_raw + 167936); + const int smem_p0_addr = smem + 167936; + uint8_t* smem_p1 = reinterpret_cast(smem_raw + 169984); + const int smem_p1_addr = smem + 169984; + int* page_offsets = reinterpret_cast(smem_raw + 172032); + const int page_offsets_addr = smem + 172032; + float* exchange0 = reinterpret_cast(smem_raw + 172800); + const int exchange0_addr = smem + 172800; + float* exchange1 = reinterpret_cast(smem_raw + 173056); + const int exchange1_addr = smem + 173056; + unsigned int* exchange0_u32 = reinterpret_cast(smem_raw + 172800); + const int exchange0_u32_addr = smem + 172800; + unsigned int* exchange1_u32 = reinterpret_cast(smem_raw + 173056); + const int exchange1_u32_addr = smem + 173056; + float* corr_stats_reduce = reinterpret_cast(smem_raw + 173312); + const int corr_stats_reduce_addr = smem + 173312; + int* reduction_ticket = reinterpret_cast(smem_raw + 173696); + const int reduction_ticket_addr = smem + 173696; + uint8_t* smem_o = reinterpret_cast(smem_raw + 173824); + const int smem_o_addr = smem + 173824; + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&Q_map))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&K_map))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&V_map))) : "memory"); + + // Mbarrier init (25 pipeline groups, 0 ordered-sequence groups, 64 barriers) + // Mbarriers at smem_raw[177920..178432) + + if (warp == 0) { + // --- pipeline 'q_pipe' --- + // q_full: 2 barriers, init_count=1 + // q_empty: 2 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 9 barriers, init_count=1 + // kv_empty: 9 barriers, init_count=1 + // s_full0: 1 barriers, init_count=1 + // s_full1: 1 barriers, init_count=1 + // _source_p_empty0: 1 barriers, init_count=1 + // _source_p_empty1: 1 barriers, init_count=1 + // o_full0: 1 barriers, init_count=1 + // o_full1: 1 barriers, init_count=1 + // --- pipeline 'page_pipe' --- + // page_full: 6 barriers, init_count=32 + // page_empty: 6 barriers, init_count=32 + // s_empty0: 1 barriers, init_count=128 + // scale_ready0: 1 barriers, init_count=128 + // stats_empty0: 1 barriers, init_count=128 + // s_empty1: 1 barriers, init_count=128 + // scale_ready1: 1 barriers, init_count=128 + // stats_empty1: 1 barriers, init_count=128 + // order_p01_0: 1 barriers, init_count=128 + // order_p01_1: 1 barriers, init_count=128 + // _source_p_ready0: 1 barriers, init_count=128 + // _source_p_ready1: 1 barriers, init_count=128 + // o_empty0: 1 barriers, init_count=128 + // o_empty1: 1 barriers, init_count=128 + // _source_barrier_reservation: 12 barriers, init_count=128 + // Warp-cooperative initialization in physical record order. + uint32_t _mbarrier_init_count_0_0 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_0) : "r"(lane), "n"(28), "r"((uint32_t)(1))); + mbarrier_init(smem + 177920 + lane * 8, _mbarrier_init_count_0_0); + uint32_t _mbarrier_init_count_0_32 = 128; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(32))); + mbarrier_init(smem + 178176 + lane * 8, _mbarrier_init_count_0_32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 178432); + if (warp == 0) { + int _tmem_hold = smem + 178432; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: softmax0 ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 176;"); + { // softmax0_main + int cta_q_s0 = blockIdx.x / 2; + int cta_kv_s0 = blockIdx.x % 2; + int batch_s0 = blockIdx.z; + int seq_len_s0 = seq_lens[batch_s0] - (q_len - 1 - cta_q_s0); + int num_ctas_s0 = 2; + int num_steps_s0 = 2; + const int warp_in_group = warp % 4; + const int wg_tid = (unsigned int)(warp_in_group * 32) + lane; + const int col_group = wg_tid % 4; + const int col_group_base = col_group * 4; + float row_max[4]; + float row_sum[4]; + row_max[0] = -3.4028235e+38f; + row_max[1] = -3.4028235e+38f; + row_max[2] = -3.4028235e+38f; + row_max[3] = -3.4028235e+38f; + row_sum[0] = 0.0f; + row_sum[1] = 0.0f; + row_sum[2] = 0.0f; + row_sum[3] = 0.0f; + uint32_t _amf_u_0 = __float_as_uint(-3.4028235e+38f); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + if (wg_tid < 16) { + exchange0_u32[wg_tid] = _amf_enc_0; + } + asm volatile("barrier.sync 8, 128;" ::: "memory"); + int order_phase = 1; + int order_stage = 0; + unsigned int _phase_s_full0_0 = 0; + unsigned int _phase_stats_empty0_0 = 1; + #pragma unroll 1 + for (int tile = 0; tile < num_steps_s0 * 2; tile += 2) { + mbarrier_wait(s_full0_addr, _phase_s_full0_0); + _phase_s_full0_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float scores[16]; + float scores_lo[8]; + float scores_hi[8]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&scores_lo[0])), "=r"(*reinterpret_cast(&scores_lo[1])), "=r"(*reinterpret_cast(&scores_lo[2])), "=r"(*reinterpret_cast(&scores_lo[3])), "=r"(*reinterpret_cast(&scores_lo[4])), "=r"(*reinterpret_cast(&scores_lo[5])), "=r"(*reinterpret_cast(&scores_lo[6])), "=r"(*reinterpret_cast(&scores_lo[7])) + : "r"(taddr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&scores_hi[0])), "=r"(*reinterpret_cast(&scores_hi[1])), "=r"(*reinterpret_cast(&scores_hi[2])), "=r"(*reinterpret_cast(&scores_hi[3])), "=r"(*reinterpret_cast(&scores_hi[4])), "=r"(*reinterpret_cast(&scores_hi[5])), "=r"(*reinterpret_cast(&scores_hi[6])), "=r"(*reinterpret_cast(&scores_hi[7])) + : "r"(taddr + 1048576)); + #pragma unroll + for (int value_idx = 0; value_idx < 8; value_idx++) { + scores[value_idx] = scores_lo[value_idx]; + scores[value_idx + 8] = scores_hi[value_idx]; + } + int tile_token_base = (cta_kv_s0 * (num_steps_s0 * 2) + tile) * 128; + int local_token_base = (unsigned int)(warp_in_group * 32) + lane / 4; + if (seq_len_s0 < tile_token_base + 128) { + if (seq_len_s0 <= tile_token_base + local_token_base) { + scores[0] = -3.4028235e+38f; + scores[1] = -3.4028235e+38f; + scores[4] = -3.4028235e+38f; + scores[5] = -3.4028235e+38f; + } + if (seq_len_s0 <= tile_token_base + local_token_base + 8) { + scores[2] = -3.4028235e+38f; + scores[3] = -3.4028235e+38f; + scores[6] = -3.4028235e+38f; + scores[7] = -3.4028235e+38f; + } + if (seq_len_s0 <= tile_token_base + local_token_base + 16) { + scores[8] = -3.4028235e+38f; + scores[9] = -3.4028235e+38f; + scores[12] = -3.4028235e+38f; + scores[13] = -3.4028235e+38f; + } + if (seq_len_s0 <= tile_token_base + local_token_base + 24) { + scores[10] = -3.4028235e+38f; + scores[11] = -3.4028235e+38f; + scores[14] = -3.4028235e+38f; + scores[15] = -3.4028235e+38f; + } + } + float tile_max[4]; + float _max_0 = max_noftz(scores[0], scores[2]); + float _max_1 = max_noftz(scores[8], scores[10]); + float _max_2 = max_noftz(_max_0, _max_1); + tile_max[0] = _max_2; + float _max_3 = max_noftz(scores[1], scores[3]); + float _max_4 = max_noftz(scores[9], scores[11]); + float _max_5 = max_noftz(_max_3, _max_4); + tile_max[1] = _max_5; + float _max_6 = max_noftz(scores[4], scores[6]); + float _max_7 = max_noftz(scores[12], scores[14]); + float _max_8 = max_noftz(_max_6, _max_7); + tile_max[2] = _max_8; + float _max_9 = max_noftz(scores[5], scores[7]); + float _max_10 = max_noftz(scores[13], scores[15]); + float _max_11 = max_noftz(_max_9, _max_10); + tile_max[3] = _max_11; + const int local_row_idx = lane / 4 % 4; + float pair01_left = ((local_row_idx % 2 == 0) ? tile_max[1] : tile_max[0]); + float pair01_right = ((local_row_idx % 2 == 0) ? tile_max[0] : tile_max[1]); + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, pair01_left, 4); + float _max_12 = max_noftz(_shfl_xor_0, pair01_right); + float pair01 = _max_12; + float pair23_left = ((local_row_idx % 2 == 0) ? tile_max[3] : tile_max[2]); + float pair23_right = ((local_row_idx % 2 == 0) ? tile_max[2] : tile_max[3]); + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, pair23_left, 4); + float _max_13 = max_noftz(_shfl_xor_1, pair23_right); + float pair23 = _max_13; + float group_left = ((local_row_idx < 2) ? pair23 : pair01); + float group_right = ((local_row_idx < 2) ? pair01 : pair23); + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, group_left, 8); + float _max_14 = max_noftz(_shfl_xor_2, group_right); + float reduced_max = _max_14; + uint32_t _amf_u_1 = __float_as_uint(reduced_max); + uint32_t _amf_mask_1 = -int32_t(_amf_u_1 >> 31) | 0x80000000u; + unsigned int _amf_enc_1 = _amf_u_1 ^ _amf_mask_1; + atomicMax(&exchange0_u32[col_group_base + local_row_idx], _amf_enc_1); + asm volatile("barrier.sync 8, 128;" ::: "memory"); + float old_max[4]; + float new_max[4]; + #pragma unroll + for (int stat = 0; stat < 4; stat++) { + old_max[stat] = row_max[stat]; + uint32_t _amf_u_2 = exchange0_u32[col_group_base + stat]; + uint32_t _amf_mask_2 = ((_amf_u_2 >> 31) - 1u) | 0x80000000u; + float _amf_dec_0 = __uint_as_float(_amf_u_2 ^ _amf_mask_2); + new_max[stat] = _amf_dec_0; + row_max[stat] = new_max[stat]; + } + mbarrier_wait(stats_empty0_addr, _phase_stats_empty0_0); + _phase_stats_empty0_0 ^= 1; + float local_stats[8]; + #pragma unroll + for (int stat_1 = 0; stat_1 < 4; stat_1++) { + local_stats[stat_1] = old_max[stat_1]; + local_stats[stat_1 + 4] = new_max[stat_1]; + } + tmem_st_x8_f32(taddr + 32, local_stats); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(scale_ready0_addr); + float safe_max0 = ((new_max[0] == -3.4028235e+38f) ? 0.0f : new_max[0]); + float safe_max1 = ((new_max[1] == -3.4028235e+38f) ? 0.0f : new_max[1]); + float safe_max2 = ((new_max[2] == -3.4028235e+38f) ? 0.0f : new_max[2]); + float safe_max3 = ((new_max[3] == -3.4028235e+38f) ? 0.0f : new_max[3]); + float2 _f2_0 = make_float2(bmm1_scale_log2, bmm1_scale_log2); + float2 _f2_1 = make_float2(-bmm1_scale_log2, -bmm1_scale_log2); + float2 _f2_2 = make_float2(8.8073549f, 8.8073549f); + float2 _f2_3 = make_float2(safe_max0, safe_max1); + float2 _f2_4 = make_float2(safe_max2, safe_max3); + float2 neg_max01 = fma_f32x2_rn_ftz(_f2_3, _f2_1, _f2_2); + float2 neg_max23 = fma_f32x2_rn_ftz(_f2_4, _f2_1, _f2_2); + float probabilities[16]; + float _fma_0 = __fmaf_rn(bmm1_scale_log2, scores[0], neg_max01.x); + probabilities[0] = _fma_0; + float _fma_1 = __fmaf_rn(bmm1_scale_log2, scores[1], neg_max01.y); + probabilities[1] = _fma_1; + float2 _f2_5 = make_float2(scores[2], scores[3]); + float2 _mul_f32x2_0; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_0) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_5)); + float2 scaled23 = _mul_f32x2_0; + float2 prob23 = add_f32x2(scaled23, neg_max01); + probabilities[2] = prob23.x; + probabilities[3] = prob23.y; + mbarrier_wait(order_p01_0_addr, order_phase); + float _exp2_0 = approx_exp2(probabilities[0]); + probabilities[0] = _exp2_0; + float2 _f2_6 = make_float2(scores[4], scores[5]); + float2 _mul_f32x2_1; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_1) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_6)); + float2 scaled45 = _mul_f32x2_1; + float2 prob45 = add_f32x2(scaled45, neg_max23); + probabilities[4] = prob45.x; + probabilities[5] = prob45.y; + float _exp2_1 = approx_exp2(probabilities[1]); + probabilities[1] = _exp2_1; + float _exp2_2 = approx_exp2(probabilities[2]); + probabilities[2] = _exp2_2; + float2 _f2_7 = make_float2(scores[6], scores[7]); + float2 _mul_f32x2_2; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_2) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_7)); + float2 scaled67 = _mul_f32x2_2; + float2 prob67 = add_f32x2(scaled67, neg_max23); + probabilities[6] = prob67.x; + probabilities[7] = prob67.y; + float _exp2_3 = approx_exp2(probabilities[3]); + probabilities[3] = _exp2_3; + float _exp2_4 = approx_exp2(probabilities[4]); + probabilities[4] = _exp2_4; + float2 _f2_8 = make_float2(scores[8], scores[9]); + float2 _mul_f32x2_3; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_3) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_8)); + float2 scaled89 = _mul_f32x2_3; + float2 prob89 = add_f32x2(scaled89, neg_max01); + probabilities[8] = prob89.x; + probabilities[9] = prob89.y; + float _exp2_5 = approx_exp2(probabilities[5]); + probabilities[5] = _exp2_5; + float _exp2_6 = approx_exp2(probabilities[6]); + probabilities[6] = _exp2_6; + float2 _f2_9 = make_float2(scores[10], scores[11]); + float2 _mul_f32x2_4; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_4) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_9)); + float2 scaled1011 = _mul_f32x2_4; + float2 prob1011 = add_f32x2(scaled1011, neg_max01); + probabilities[10] = prob1011.x; + probabilities[11] = prob1011.y; + float _exp2_7 = approx_exp2(probabilities[7]); + probabilities[7] = _exp2_7; + float _exp2_8 = approx_exp2(probabilities[8]); + probabilities[8] = _exp2_8; + float2 _f2_10 = make_float2(scores[12], scores[13]); + float2 _mul_f32x2_5; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_5) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_10)); + float2 scaled1213 = _mul_f32x2_5; + float2 prob1213 = add_f32x2(scaled1213, neg_max23); + probabilities[12] = prob1213.x; + probabilities[13] = prob1213.y; + float _exp2_9 = approx_exp2(probabilities[9]); + probabilities[9] = _exp2_9; + mbarrier_arrive(order_p01_1_addr); + order_phase ^= 1; + float _exp2_10 = approx_exp2(probabilities[10]); + probabilities[10] = _exp2_10; + float2 _f2_11 = make_float2(scores[14], scores[15]); + float2 _mul_f32x2_6; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_6) : "l"(*(const unsigned long long*)&_f2_0), "l"(*(const unsigned long long*)&_f2_11)); + float2 scaled1415 = _mul_f32x2_6; + float2 prob1415 = add_f32x2(scaled1415, neg_max23); + probabilities[14] = prob1415.x; + probabilities[15] = prob1415.y; + float _exp2_11 = approx_exp2(probabilities[11]); + probabilities[11] = _exp2_11; + unsigned int regs_p[4]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities[0]), "f"(probabilities[1]), + "f"(probabilities[2]), "f"(probabilities[3])); + regs_p[0] = _packed; + } + #pragma unroll + for (int value_idx_1 = 12; value_idx_1 < 16; value_idx_1++) { + float _exp2_12 = approx_exp2(probabilities[value_idx_1]); + probabilities[value_idx_1] = _exp2_12; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities[4]), "f"(probabilities[5]), + "f"(probabilities[6]), "f"(probabilities[7])); + regs_p[1] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities[8]), "f"(probabilities[9]), + "f"(probabilities[10]), "f"(probabilities[11])); + regs_p[2] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities[12]), "f"(probabilities[13]), + "f"(probabilities[14]), "f"(probabilities[15])); + regs_p[3] = _packed; + } + const int matrix_idx = lane / 8; + const int matrix_row_idx = matrix_idx % 2; + const int matrix_col_idx = matrix_idx / 2; + const int thread_row_idx = lane % 8; + const int segment_col_idx = warp_in_group * 2 + matrix_col_idx ^ thread_row_idx; + const int stsm_offset = (matrix_row_idx * 8 + thread_row_idx) * 128 + segment_col_idx * 16; + const void* _stmatrix_b8_ptr_3 = reinterpret_cast(reinterpret_cast(smem_p0) + stsm_offset); + uint64_t _stmatrix_b8_addr64_3; + asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(_stmatrix_b8_addr64_3) : "l"(_stmatrix_b8_ptr_3)); + uint32_t _stmatrix_b8_addr_3; + asm volatile("cvt.u32.u64 %0, %1;" : "=r"(_stmatrix_b8_addr_3) : "l"(_stmatrix_b8_addr64_3)); + asm volatile("stmatrix.sync.aligned.m16n8.x4.trans.shared.b8 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_b8_addr_3), "r"(regs_p[0]), "r"(regs_p[1]), "r"(regs_p[2]), "r"(regs_p[3]) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (tile + 2 >= num_steps_s0 * 2) { + mbarrier_wait(stats_empty0_addr, _phase_stats_empty0_0); + _phase_stats_empty0_0 ^= 1; + } + mbarrier_arrive(s_empty0_addr); + float acc_scale[4]; + #pragma unroll + for (int stat_2 = 0; stat_2 < 4; stat_2++) { + float max_delta = old_max[stat_2] - new_max[stat_2]; + acc_scale[stat_2] = 1.0f; + if (max_delta != 0.0f) { + float _exp2_13 = approx_exp2(max_delta * bmm1_scale_log2); + acc_scale[stat_2] = _exp2_13; + } + } + float2 _f2_12 = make_float2(row_sum[0], row_sum[1]); + float2 _f2_13 = make_float2(acc_scale[0], acc_scale[1]); + float2 _f2_14 = make_float2(probabilities[0], probabilities[1]); + float2 _f2_15 = make_float2(probabilities[2], probabilities[3]); + float2 _f2_16 = make_float2(probabilities[8], probabilities[9]); + float2 _f2_17 = make_float2(probabilities[10], probabilities[11]); + float2 next_sum01 = fma_f32x2_rn_ftz(_f2_12, _f2_13, _f2_14); + next_sum01 = add_f32x2(next_sum01, _f2_15); + next_sum01 = add_f32x2(next_sum01, _f2_16); + next_sum01 = add_f32x2(next_sum01, _f2_17); + row_sum[0] = next_sum01.x; + row_sum[1] = next_sum01.y; + float2 _f2_18 = make_float2(row_sum[2], row_sum[3]); + float2 _f2_19 = make_float2(acc_scale[2], acc_scale[3]); + float2 _f2_20 = make_float2(probabilities[4], probabilities[5]); + float2 _f2_21 = make_float2(probabilities[6], probabilities[7]); + float2 _f2_22 = make_float2(probabilities[12], probabilities[13]); + float2 _f2_23 = make_float2(probabilities[14], probabilities[15]); + float2 next_sum23 = fma_f32x2_rn_ftz(_f2_18, _f2_19, _f2_20); + next_sum23 = add_f32x2(next_sum23, _f2_21); + next_sum23 = add_f32x2(next_sum23, _f2_22); + next_sum23 = add_f32x2(next_sum23, _f2_23); + row_sum[2] = next_sum23.x; + row_sum[3] = next_sum23.y; + if (tile + 2 >= num_steps_s0 * 2) { + float final_stats[8]; + #pragma unroll + for (int stat_3 = 0; stat_3 < 4; stat_3++) { + final_stats[stat_3] = row_sum[stat_3]; + final_stats[stat_3 + 4] = row_max[stat_3]; + } + tmem_st_x8_f32(taddr + 32, final_stats); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(scale_ready0_addr); + } + } + } + // ---- Role: softmax1 ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 176;"); + { // softmax1_main + int cta_q_s1 = blockIdx.x / 2; + int cta_kv_s1 = blockIdx.x % 2; + int batch_s1 = blockIdx.z; + int seq_len_s1 = seq_lens[batch_s1] - (q_len - 1 - cta_q_s1); + int num_ctas_s1 = 2; + int num_steps_s1 = 2; + const int warp_in_group_1 = warp % 4; + const int wg_tid_1 = (unsigned int)(warp_in_group_1 * 32) + lane; + const int col_group_1 = wg_tid_1 % 4; + const int col_group_base_1 = col_group_1 * 4; + float row_max_1[4]; + float row_sum_1[4]; + row_max_1[0] = -3.4028235e+38f; + row_max_1[1] = -3.4028235e+38f; + row_max_1[2] = -3.4028235e+38f; + row_max_1[3] = -3.4028235e+38f; + row_sum_1[0] = 0.0f; + row_sum_1[1] = 0.0f; + row_sum_1[2] = 0.0f; + row_sum_1[3] = 0.0f; + uint32_t _amf_u_0 = __float_as_uint(-3.4028235e+38f); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_2 = _amf_u_0 ^ _amf_mask_0; + if (wg_tid_1 < 16) { + exchange1_u32[wg_tid_1] = _amf_enc_2; + } + asm volatile("barrier.sync 9, 128;" ::: "memory"); + int order_phase_1 = 0; + int order_stage_1 = 0; + unsigned int _phase_s_full1_0 = 0; + unsigned int _phase_stats_empty1_0 = 1; + #pragma unroll 1 + for (int tile_1 = 1; tile_1 < num_steps_s1 * 2; tile_1 += 2) { + mbarrier_wait(s_full1_addr, _phase_s_full1_0); + _phase_s_full1_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float scores_1[16]; + float scores_lo_1[8]; + float scores_hi_1[8]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&scores_lo_1[0])), "=r"(*reinterpret_cast(&scores_lo_1[1])), "=r"(*reinterpret_cast(&scores_lo_1[2])), "=r"(*reinterpret_cast(&scores_lo_1[3])), "=r"(*reinterpret_cast(&scores_lo_1[4])), "=r"(*reinterpret_cast(&scores_lo_1[5])), "=r"(*reinterpret_cast(&scores_lo_1[6])), "=r"(*reinterpret_cast(&scores_lo_1[7])) + : "r"(taddr + 16)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&scores_hi_1[0])), "=r"(*reinterpret_cast(&scores_hi_1[1])), "=r"(*reinterpret_cast(&scores_hi_1[2])), "=r"(*reinterpret_cast(&scores_hi_1[3])), "=r"(*reinterpret_cast(&scores_hi_1[4])), "=r"(*reinterpret_cast(&scores_hi_1[5])), "=r"(*reinterpret_cast(&scores_hi_1[6])), "=r"(*reinterpret_cast(&scores_hi_1[7])) + : "r"(taddr + 16 + 1048576)); + #pragma unroll + for (int value_idx_2 = 0; value_idx_2 < 8; value_idx_2++) { + scores_1[value_idx_2] = scores_lo_1[value_idx_2]; + scores_1[value_idx_2 + 8] = scores_hi_1[value_idx_2]; + } + int tile_token_base_1 = (cta_kv_s1 * (num_steps_s1 * 2) + tile_1) * 128; + int local_token_base_1 = (unsigned int)(warp_in_group_1 * 32) + lane / 4; + if (seq_len_s1 < tile_token_base_1 + 128) { + if (seq_len_s1 <= tile_token_base_1 + local_token_base_1) { + scores_1[0] = -3.4028235e+38f; + scores_1[1] = -3.4028235e+38f; + scores_1[4] = -3.4028235e+38f; + scores_1[5] = -3.4028235e+38f; + } + if (seq_len_s1 <= tile_token_base_1 + local_token_base_1 + 8) { + scores_1[2] = -3.4028235e+38f; + scores_1[3] = -3.4028235e+38f; + scores_1[6] = -3.4028235e+38f; + scores_1[7] = -3.4028235e+38f; + } + if (seq_len_s1 <= tile_token_base_1 + local_token_base_1 + 16) { + scores_1[8] = -3.4028235e+38f; + scores_1[9] = -3.4028235e+38f; + scores_1[12] = -3.4028235e+38f; + scores_1[13] = -3.4028235e+38f; + } + if (seq_len_s1 <= tile_token_base_1 + local_token_base_1 + 24) { + scores_1[10] = -3.4028235e+38f; + scores_1[11] = -3.4028235e+38f; + scores_1[14] = -3.4028235e+38f; + scores_1[15] = -3.4028235e+38f; + } + } + float tile_max_1[4]; + float _max_15 = max_noftz(scores_1[0], scores_1[2]); + float _max_16 = max_noftz(scores_1[8], scores_1[10]); + float _max_17 = max_noftz(_max_15, _max_16); + tile_max_1[0] = _max_17; + float _max_18 = max_noftz(scores_1[1], scores_1[3]); + float _max_19 = max_noftz(scores_1[9], scores_1[11]); + float _max_20 = max_noftz(_max_18, _max_19); + tile_max_1[1] = _max_20; + float _max_21 = max_noftz(scores_1[4], scores_1[6]); + float _max_22 = max_noftz(scores_1[12], scores_1[14]); + float _max_23 = max_noftz(_max_21, _max_22); + tile_max_1[2] = _max_23; + float _max_24 = max_noftz(scores_1[5], scores_1[7]); + float _max_25 = max_noftz(scores_1[13], scores_1[15]); + float _max_26 = max_noftz(_max_24, _max_25); + tile_max_1[3] = _max_26; + const int local_row_idx_1 = lane / 4 % 4; + float pair01_left_1 = ((local_row_idx_1 % 2 == 0) ? tile_max_1[1] : tile_max_1[0]); + float pair01_right_1 = ((local_row_idx_1 % 2 == 0) ? tile_max_1[0] : tile_max_1[1]); + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, pair01_left_1, 4); + float _max_27 = max_noftz(_shfl_xor_3, pair01_right_1); + float pair01_1 = _max_27; + float pair23_left_1 = ((local_row_idx_1 % 2 == 0) ? tile_max_1[3] : tile_max_1[2]); + float pair23_right_1 = ((local_row_idx_1 % 2 == 0) ? tile_max_1[2] : tile_max_1[3]); + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, pair23_left_1, 4); + float _max_28 = max_noftz(_shfl_xor_4, pair23_right_1); + float pair23_1 = _max_28; + float group_left_1 = ((local_row_idx_1 < 2) ? pair23_1 : pair01_1); + float group_right_1 = ((local_row_idx_1 < 2) ? pair01_1 : pair23_1); + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, group_left_1, 8); + float _max_29 = max_noftz(_shfl_xor_5, group_right_1); + float reduced_max_1 = _max_29; + uint32_t _amf_u_1 = __float_as_uint(reduced_max_1); + uint32_t _amf_mask_1 = -int32_t(_amf_u_1 >> 31) | 0x80000000u; + unsigned int _amf_enc_3 = _amf_u_1 ^ _amf_mask_1; + atomicMax(&exchange1_u32[col_group_base_1 + local_row_idx_1], _amf_enc_3); + asm volatile("barrier.sync 9, 128;" ::: "memory"); + float old_max_1[4]; + float new_max_1[4]; + #pragma unroll + for (int stat_4 = 0; stat_4 < 4; stat_4++) { + old_max_1[stat_4] = row_max_1[stat_4]; + uint32_t _amf_u_2 = exchange1_u32[col_group_base_1 + stat_4]; + uint32_t _amf_mask_2 = ((_amf_u_2 >> 31) - 1u) | 0x80000000u; + float _amf_dec_1 = __uint_as_float(_amf_u_2 ^ _amf_mask_2); + new_max_1[stat_4] = _amf_dec_1; + row_max_1[stat_4] = new_max_1[stat_4]; + } + mbarrier_wait(stats_empty1_addr, _phase_stats_empty1_0); + _phase_stats_empty1_0 ^= 1; + float local_stats_1[8]; + #pragma unroll + for (int stat_5 = 0; stat_5 < 4; stat_5++) { + local_stats_1[stat_5] = old_max_1[stat_5]; + local_stats_1[stat_5 + 4] = new_max_1[stat_5]; + } + tmem_st_x8_f32(taddr + 64, local_stats_1); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(scale_ready1_addr); + float safe_max0_1 = ((new_max_1[0] == -3.4028235e+38f) ? 0.0f : new_max_1[0]); + float safe_max1_1 = ((new_max_1[1] == -3.4028235e+38f) ? 0.0f : new_max_1[1]); + float safe_max2_1 = ((new_max_1[2] == -3.4028235e+38f) ? 0.0f : new_max_1[2]); + float safe_max3_1 = ((new_max_1[3] == -3.4028235e+38f) ? 0.0f : new_max_1[3]); + float2 _f2_24 = make_float2(bmm1_scale_log2, bmm1_scale_log2); + float2 _f2_25 = make_float2(-bmm1_scale_log2, -bmm1_scale_log2); + float2 _f2_26 = make_float2(8.8073549f, 8.8073549f); + float2 _f2_27 = make_float2(safe_max0_1, safe_max1_1); + float2 _f2_28 = make_float2(safe_max2_1, safe_max3_1); + float2 neg_max01_1 = fma_f32x2_rn_ftz(_f2_27, _f2_25, _f2_26); + float2 neg_max23_1 = fma_f32x2_rn_ftz(_f2_28, _f2_25, _f2_26); + float probabilities_1[16]; + float _fma_2 = __fmaf_rn(bmm1_scale_log2, scores_1[0], neg_max01_1.x); + probabilities_1[0] = _fma_2; + float _fma_3 = __fmaf_rn(bmm1_scale_log2, scores_1[1], neg_max01_1.y); + probabilities_1[1] = _fma_3; + float2 _f2_29 = make_float2(scores_1[2], scores_1[3]); + float2 _mul_f32x2_7; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_7) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_29)); + float2 scaled23_1 = _mul_f32x2_7; + float2 prob23_1 = add_f32x2(scaled23_1, neg_max01_1); + probabilities_1[2] = prob23_1.x; + probabilities_1[3] = prob23_1.y; + mbarrier_wait(order_p01_1_addr, order_phase_1); + float _exp2_14 = approx_exp2(probabilities_1[0]); + probabilities_1[0] = _exp2_14; + float2 _f2_30 = make_float2(scores_1[4], scores_1[5]); + float2 _mul_f32x2_8; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_8) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_30)); + float2 scaled45_1 = _mul_f32x2_8; + float2 prob45_1 = add_f32x2(scaled45_1, neg_max23_1); + probabilities_1[4] = prob45_1.x; + probabilities_1[5] = prob45_1.y; + float _exp2_15 = approx_exp2(probabilities_1[1]); + probabilities_1[1] = _exp2_15; + float _exp2_16 = approx_exp2(probabilities_1[2]); + probabilities_1[2] = _exp2_16; + float2 _f2_31 = make_float2(scores_1[6], scores_1[7]); + float2 _mul_f32x2_9; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_9) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_31)); + float2 scaled67_1 = _mul_f32x2_9; + float2 prob67_1 = add_f32x2(scaled67_1, neg_max23_1); + probabilities_1[6] = prob67_1.x; + probabilities_1[7] = prob67_1.y; + float _exp2_17 = approx_exp2(probabilities_1[3]); + probabilities_1[3] = _exp2_17; + float _exp2_18 = approx_exp2(probabilities_1[4]); + probabilities_1[4] = _exp2_18; + float2 _f2_32 = make_float2(scores_1[8], scores_1[9]); + float2 _mul_f32x2_10; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_10) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_32)); + float2 scaled89_1 = _mul_f32x2_10; + float2 prob89_1 = add_f32x2(scaled89_1, neg_max01_1); + probabilities_1[8] = prob89_1.x; + probabilities_1[9] = prob89_1.y; + float _exp2_19 = approx_exp2(probabilities_1[5]); + probabilities_1[5] = _exp2_19; + float _exp2_20 = approx_exp2(probabilities_1[6]); + probabilities_1[6] = _exp2_20; + float2 _f2_33 = make_float2(scores_1[10], scores_1[11]); + float2 _mul_f32x2_11; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_11) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_33)); + float2 scaled1011_1 = _mul_f32x2_11; + float2 prob1011_1 = add_f32x2(scaled1011_1, neg_max01_1); + probabilities_1[10] = prob1011_1.x; + probabilities_1[11] = prob1011_1.y; + float _exp2_21 = approx_exp2(probabilities_1[7]); + probabilities_1[7] = _exp2_21; + float _exp2_22 = approx_exp2(probabilities_1[8]); + probabilities_1[8] = _exp2_22; + float2 _f2_34 = make_float2(scores_1[12], scores_1[13]); + float2 _mul_f32x2_12; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_12) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_34)); + float2 scaled1213_1 = _mul_f32x2_12; + float2 prob1213_1 = add_f32x2(scaled1213_1, neg_max23_1); + probabilities_1[12] = prob1213_1.x; + probabilities_1[13] = prob1213_1.y; + float _exp2_23 = approx_exp2(probabilities_1[9]); + probabilities_1[9] = _exp2_23; + mbarrier_arrive(order_p01_0_addr); + order_phase_1 ^= 1; + float _exp2_24 = approx_exp2(probabilities_1[10]); + probabilities_1[10] = _exp2_24; + float2 _f2_35 = make_float2(scores_1[14], scores_1[15]); + float2 _mul_f32x2_13; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_13) : "l"(*(const unsigned long long*)&_f2_24), "l"(*(const unsigned long long*)&_f2_35)); + float2 scaled1415_1 = _mul_f32x2_13; + float2 prob1415_1 = add_f32x2(scaled1415_1, neg_max23_1); + probabilities_1[14] = prob1415_1.x; + probabilities_1[15] = prob1415_1.y; + float _exp2_25 = approx_exp2(probabilities_1[11]); + probabilities_1[11] = _exp2_25; + unsigned int regs_p_1[4]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities_1[0]), "f"(probabilities_1[1]), + "f"(probabilities_1[2]), "f"(probabilities_1[3])); + regs_p_1[0] = _packed; + } + #pragma unroll + for (int value_idx_3 = 12; value_idx_3 < 16; value_idx_3++) { + float _exp2_26 = approx_exp2(probabilities_1[value_idx_3]); + probabilities_1[value_idx_3] = _exp2_26; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities_1[4]), "f"(probabilities_1[5]), + "f"(probabilities_1[6]), "f"(probabilities_1[7])); + regs_p_1[1] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities_1[8]), "f"(probabilities_1[9]), + "f"(probabilities_1[10]), "f"(probabilities_1[11])); + regs_p_1[2] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(probabilities_1[12]), "f"(probabilities_1[13]), + "f"(probabilities_1[14]), "f"(probabilities_1[15])); + regs_p_1[3] = _packed; + } + const int matrix_idx_1 = lane / 8; + const int matrix_row_idx_1 = matrix_idx_1 % 2; + const int matrix_col_idx_1 = matrix_idx_1 / 2; + const int thread_row_idx_1 = lane % 8; + const int segment_col_idx_1 = warp_in_group_1 * 2 + matrix_col_idx_1 ^ thread_row_idx_1; + const int stsm_offset_1 = (matrix_row_idx_1 * 8 + thread_row_idx_1) * 128 + segment_col_idx_1 * 16; + const void* _stmatrix_b8_ptr_3 = reinterpret_cast(reinterpret_cast(smem_p1) + stsm_offset_1); + uint64_t _stmatrix_b8_addr64_3; + asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(_stmatrix_b8_addr64_3) : "l"(_stmatrix_b8_ptr_3)); + uint32_t _stmatrix_b8_addr_3; + asm volatile("cvt.u32.u64 %0, %1;" : "=r"(_stmatrix_b8_addr_3) : "l"(_stmatrix_b8_addr64_3)); + asm volatile("stmatrix.sync.aligned.m16n8.x4.trans.shared.b8 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_b8_addr_3), "r"(regs_p_1[0]), "r"(regs_p_1[1]), "r"(regs_p_1[2]), "r"(regs_p_1[3]) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (tile_1 + 2 >= num_steps_s1 * 2) { + mbarrier_wait(stats_empty1_addr, _phase_stats_empty1_0); + _phase_stats_empty1_0 ^= 1; + } + mbarrier_arrive(s_empty1_addr); + float acc_scale_1[4]; + #pragma unroll + for (int stat_6 = 0; stat_6 < 4; stat_6++) { + float max_delta_1 = old_max_1[stat_6] - new_max_1[stat_6]; + acc_scale_1[stat_6] = 1.0f; + if (max_delta_1 != 0.0f) { + float _exp2_27 = approx_exp2(max_delta_1 * bmm1_scale_log2); + acc_scale_1[stat_6] = _exp2_27; + } + } + float2 _f2_36 = make_float2(row_sum_1[0], row_sum_1[1]); + float2 _f2_37 = make_float2(acc_scale_1[0], acc_scale_1[1]); + float2 _f2_38 = make_float2(probabilities_1[0], probabilities_1[1]); + float2 _f2_39 = make_float2(probabilities_1[2], probabilities_1[3]); + float2 _f2_40 = make_float2(probabilities_1[8], probabilities_1[9]); + float2 _f2_41 = make_float2(probabilities_1[10], probabilities_1[11]); + float2 next_sum01_1 = fma_f32x2_rn_ftz(_f2_36, _f2_37, _f2_38); + next_sum01_1 = add_f32x2(next_sum01_1, _f2_39); + next_sum01_1 = add_f32x2(next_sum01_1, _f2_40); + next_sum01_1 = add_f32x2(next_sum01_1, _f2_41); + row_sum_1[0] = next_sum01_1.x; + row_sum_1[1] = next_sum01_1.y; + float2 _f2_42 = make_float2(row_sum_1[2], row_sum_1[3]); + float2 _f2_43 = make_float2(acc_scale_1[2], acc_scale_1[3]); + float2 _f2_44 = make_float2(probabilities_1[4], probabilities_1[5]); + float2 _f2_45 = make_float2(probabilities_1[6], probabilities_1[7]); + float2 _f2_46 = make_float2(probabilities_1[12], probabilities_1[13]); + float2 _f2_47 = make_float2(probabilities_1[14], probabilities_1[15]); + float2 next_sum23_1 = fma_f32x2_rn_ftz(_f2_42, _f2_43, _f2_44); + next_sum23_1 = add_f32x2(next_sum23_1, _f2_45); + next_sum23_1 = add_f32x2(next_sum23_1, _f2_46); + next_sum23_1 = add_f32x2(next_sum23_1, _f2_47); + row_sum_1[2] = next_sum23_1.x; + row_sum_1[3] = next_sum23_1.y; + if (tile_1 + 2 >= num_steps_s1 * 2) { + float final_stats_1[8]; + #pragma unroll + for (int stat_7 = 0; stat_7 < 4; stat_7++) { + final_stats_1[stat_7] = row_sum_1[stat_7]; + final_stats_1[stat_7 + 4] = row_max_1[stat_7]; + } + tmem_st_x8_f32(taddr + 64, final_stats_1); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(scale_ready1_addr); + } + } + } + // ---- Role: correction ---- + } else if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + { // correction_main + int cta_q_c = blockIdx.x / 2; + int cta_kv_c = blockIdx.x % 2; + int y_c = blockIdx.y; + int head_group_c = y_c / 4; + int hv_partition_c = y_c % 4; + int batch_c = blockIdx.z; + int seq_len_c = seq_lens[batch_c] - (q_len - 1 - cta_q_c); + int num_ctas_c = 2; + int num_steps_c = 2; + int num_tiles_c = num_steps_c * 2; + const int warp_c = warp % 4; + const int corr_tid = (unsigned int)(warp_c * 32) + lane; + const int tmem_row_c = warp_c * 32 << 16; + const int corr_head_base_c = corr_tid % 4 * 4; + unsigned int _phase_scale_ready0_0 = 0; + unsigned int _phase_o_full0_0 = 0; + unsigned int _phase_scale_ready1_0 = 0; + unsigned int _phase_o_full1_0 = 0; + #pragma unroll 1 + for (int tile_c = 0; tile_c < num_tiles_c; tile_c++) { + if (tile_c % 2 == 0) { + mbarrier_wait(scale_ready0_addr, _phase_scale_ready0_0); + _phase_scale_ready0_0 ^= 1; + float corr_scale0[4]; + if (tile_c > 0) { + float _tmem_load_0[8]; + tmem_ld_x8(&_tmem_load_0[0], taddr + 32); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int head_local_c = 0; head_local_c < 4; head_local_c++) { + float corr_delta0_c = _tmem_load_0[head_local_c] - _tmem_load_0[head_local_c + 4]; + float _exp2_28 = approx_exp2(corr_delta0_c * bmm1_scale_log2); + corr_scale0[head_local_c] = ((corr_delta0_c != 0.0f) ? _exp2_28 : 1.0f); + } + } + mbarrier_arrive(stats_empty0_addr); + if (tile_c > 0) { + mbarrier_wait(o_full0_addr, _phase_o_full0_0); + _phase_o_full0_0 ^= 1; + float previous0_lo[8]; + float previous0_hi[8]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&previous0_lo[0])), "=r"(*reinterpret_cast(&previous0_lo[1])), "=r"(*reinterpret_cast(&previous0_lo[2])), "=r"(*reinterpret_cast(&previous0_lo[3])), "=r"(*reinterpret_cast(&previous0_lo[4])), "=r"(*reinterpret_cast(&previous0_lo[5])), "=r"(*reinterpret_cast(&previous0_lo[6])), "=r"(*reinterpret_cast(&previous0_lo[7])) + : "r"(taddr + 96)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&previous0_hi[0])), "=r"(*reinterpret_cast(&previous0_hi[1])), "=r"(*reinterpret_cast(&previous0_hi[2])), "=r"(*reinterpret_cast(&previous0_hi[3])), "=r"(*reinterpret_cast(&previous0_hi[4])), "=r"(*reinterpret_cast(&previous0_hi[5])), "=r"(*reinterpret_cast(&previous0_hi[6])), "=r"(*reinterpret_cast(&previous0_hi[7])) + : "r"(taddr + 96 + 1048576)); + float2 _f2_48 = make_float2(corr_scale0[0], corr_scale0[1]); + float2 _f2_49 = make_float2(corr_scale0[2], corr_scale0[3]); + int corr_pred0 = 0; + #pragma unroll + for (int head_local_c_1 = 0; head_local_c_1 < 4; head_local_c_1++) { + if (corr_scale0[head_local_c_1] != 1.0f) { + corr_pred0 = 1; + } + } + int _vote_0 = __any_sync(0xFFFFFFFF, corr_pred0 != 0); + if (_vote_0 != 0) { + #pragma unroll + for (int reg_base_c = 0; reg_base_c < 8; reg_base_c += 4) { + float2 _f2_50 = make_float2(previous0_lo[reg_base_c], previous0_lo[reg_base_c + 1]); + float2 _f2_51 = make_float2(previous0_lo[reg_base_c + 2], previous0_lo[reg_base_c + 3]); + float2 _f2_52 = make_float2(previous0_hi[reg_base_c], previous0_hi[reg_base_c + 1]); + float2 _f2_53 = make_float2(previous0_hi[reg_base_c + 2], previous0_hi[reg_base_c + 3]); + float2 _f32x2_mul_b_0 = ((reg_base_c == 0) ? _f2_48 : _f2_49); + float2 _mul_f32x2_14; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_14) : "l"(*(const unsigned long long*)&_f2_50), "l"(*(const unsigned long long*)&_f32x2_mul_b_0)); + float2 prev0_lo01_scaled_f2 = _mul_f32x2_14; + float2 _f32x2_mul_b_1 = ((reg_base_c == 0) ? _f2_48 : _f2_49); + float2 _mul_f32x2_15; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_15) : "l"(*(const unsigned long long*)&_f2_51), "l"(*(const unsigned long long*)&_f32x2_mul_b_1)); + float2 prev0_lo23_scaled_f2 = _mul_f32x2_15; + float2 _f32x2_mul_b_2 = ((reg_base_c == 0) ? _f2_48 : _f2_49); + float2 _mul_f32x2_16; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_16) : "l"(*(const unsigned long long*)&_f2_52), "l"(*(const unsigned long long*)&_f32x2_mul_b_2)); + float2 prev0_hi01_scaled_f2 = _mul_f32x2_16; + float2 _f32x2_mul_b_3 = ((reg_base_c == 0) ? _f2_48 : _f2_49); + float2 _mul_f32x2_17; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_17) : "l"(*(const unsigned long long*)&_f2_53), "l"(*(const unsigned long long*)&_f32x2_mul_b_3)); + float2 prev0_hi23_scaled_f2 = _mul_f32x2_17; + previous0_lo[reg_base_c] = prev0_lo01_scaled_f2.x; + previous0_lo[reg_base_c + 1] = prev0_lo01_scaled_f2.y; + previous0_lo[reg_base_c + 2] = prev0_lo23_scaled_f2.x; + previous0_lo[reg_base_c + 3] = prev0_lo23_scaled_f2.y; + previous0_hi[reg_base_c] = prev0_hi01_scaled_f2.x; + previous0_hi[reg_base_c + 1] = prev0_hi01_scaled_f2.y; + previous0_hi[reg_base_c + 2] = prev0_hi23_scaled_f2.x; + previous0_hi[reg_base_c + 3] = prev0_hi23_scaled_f2.y; + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x2.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(taddr + 96), "r"(*reinterpret_cast(&previous0_lo[0])), "r"(*reinterpret_cast(&previous0_lo[1])), "r"(*reinterpret_cast(&previous0_lo[2])), "r"(*reinterpret_cast(&previous0_lo[3])), "r"(*reinterpret_cast(&previous0_lo[4])), "r"(*reinterpret_cast(&previous0_lo[5])), "r"(*reinterpret_cast(&previous0_lo[6])), "r"(*reinterpret_cast(&previous0_lo[7]))); + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x2.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(taddr + 96 + 1048576), "r"(*reinterpret_cast(&previous0_hi[0])), "r"(*reinterpret_cast(&previous0_hi[1])), "r"(*reinterpret_cast(&previous0_hi[2])), "r"(*reinterpret_cast(&previous0_hi[3])), "r"(*reinterpret_cast(&previous0_hi[4])), "r"(*reinterpret_cast(&previous0_hi[5])), "r"(*reinterpret_cast(&previous0_hi[6])), "r"(*reinterpret_cast(&previous0_hi[7]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + mbarrier_arrive(o_empty0_addr); + } + } else { + mbarrier_wait(scale_ready1_addr, _phase_scale_ready1_0); + _phase_scale_ready1_0 ^= 1; + float corr_scale1[4]; + if (tile_c > 1) { + float _tmem_load_1[8]; + tmem_ld_x8(&_tmem_load_1[0], taddr + 64); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int head_local_c_2 = 0; head_local_c_2 < 4; head_local_c_2++) { + float corr_delta1_c = _tmem_load_1[head_local_c_2] - _tmem_load_1[head_local_c_2 + 4]; + float _exp2_29 = approx_exp2(corr_delta1_c * bmm1_scale_log2); + corr_scale1[head_local_c_2] = ((corr_delta1_c != 0.0f) ? _exp2_29 : 1.0f); + } + } + mbarrier_arrive(stats_empty1_addr); + if (tile_c > 1) { + mbarrier_wait(o_full1_addr, _phase_o_full1_0); + _phase_o_full1_0 ^= 1; + float previous1_lo[8]; + float previous1_hi[8]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&previous1_lo[0])), "=r"(*reinterpret_cast(&previous1_lo[1])), "=r"(*reinterpret_cast(&previous1_lo[2])), "=r"(*reinterpret_cast(&previous1_lo[3])), "=r"(*reinterpret_cast(&previous1_lo[4])), "=r"(*reinterpret_cast(&previous1_lo[5])), "=r"(*reinterpret_cast(&previous1_lo[6])), "=r"(*reinterpret_cast(&previous1_lo[7])) + : "r"(taddr + 112)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&previous1_hi[0])), "=r"(*reinterpret_cast(&previous1_hi[1])), "=r"(*reinterpret_cast(&previous1_hi[2])), "=r"(*reinterpret_cast(&previous1_hi[3])), "=r"(*reinterpret_cast(&previous1_hi[4])), "=r"(*reinterpret_cast(&previous1_hi[5])), "=r"(*reinterpret_cast(&previous1_hi[6])), "=r"(*reinterpret_cast(&previous1_hi[7])) + : "r"(taddr + 112 + 1048576)); + float2 _f2_54 = make_float2(corr_scale1[0], corr_scale1[1]); + float2 _f2_55 = make_float2(corr_scale1[2], corr_scale1[3]); + int corr_pred1 = 0; + #pragma unroll + for (int head_local_c_3 = 0; head_local_c_3 < 4; head_local_c_3++) { + if (corr_scale1[head_local_c_3] != 1.0f) { + corr_pred1 = 1; + } + } + int _vote_1 = __any_sync(0xFFFFFFFF, corr_pred1 != 0); + if (_vote_1 != 0) { + #pragma unroll + for (int reg_base_c_1 = 0; reg_base_c_1 < 8; reg_base_c_1 += 4) { + float2 _f2_56 = make_float2(previous1_lo[reg_base_c_1], previous1_lo[reg_base_c_1 + 1]); + float2 _f2_57 = make_float2(previous1_lo[reg_base_c_1 + 2], previous1_lo[reg_base_c_1 + 3]); + float2 _f2_58 = make_float2(previous1_hi[reg_base_c_1], previous1_hi[reg_base_c_1 + 1]); + float2 _f2_59 = make_float2(previous1_hi[reg_base_c_1 + 2], previous1_hi[reg_base_c_1 + 3]); + float2 _f32x2_mul_b_4 = ((reg_base_c_1 == 0) ? _f2_54 : _f2_55); + float2 _mul_f32x2_18; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_18) : "l"(*(const unsigned long long*)&_f2_56), "l"(*(const unsigned long long*)&_f32x2_mul_b_4)); + float2 prev1_lo01_scaled_f2 = _mul_f32x2_18; + float2 _f32x2_mul_b_5 = ((reg_base_c_1 == 0) ? _f2_54 : _f2_55); + float2 _mul_f32x2_19; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_19) : "l"(*(const unsigned long long*)&_f2_57), "l"(*(const unsigned long long*)&_f32x2_mul_b_5)); + float2 prev1_lo23_scaled_f2 = _mul_f32x2_19; + float2 _f32x2_mul_b_6 = ((reg_base_c_1 == 0) ? _f2_54 : _f2_55); + float2 _mul_f32x2_20; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_20) : "l"(*(const unsigned long long*)&_f2_58), "l"(*(const unsigned long long*)&_f32x2_mul_b_6)); + float2 prev1_hi01_scaled_f2 = _mul_f32x2_20; + float2 _f32x2_mul_b_7 = ((reg_base_c_1 == 0) ? _f2_54 : _f2_55); + float2 _mul_f32x2_21; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_21) : "l"(*(const unsigned long long*)&_f2_59), "l"(*(const unsigned long long*)&_f32x2_mul_b_7)); + float2 prev1_hi23_scaled_f2 = _mul_f32x2_21; + previous1_lo[reg_base_c_1] = prev1_lo01_scaled_f2.x; + previous1_lo[reg_base_c_1 + 1] = prev1_lo01_scaled_f2.y; + previous1_lo[reg_base_c_1 + 2] = prev1_lo23_scaled_f2.x; + previous1_lo[reg_base_c_1 + 3] = prev1_lo23_scaled_f2.y; + previous1_hi[reg_base_c_1] = prev1_hi01_scaled_f2.x; + previous1_hi[reg_base_c_1 + 1] = prev1_hi01_scaled_f2.y; + previous1_hi[reg_base_c_1 + 2] = prev1_hi23_scaled_f2.x; + previous1_hi[reg_base_c_1 + 3] = prev1_hi23_scaled_f2.y; + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x2.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(taddr + 112), "r"(*reinterpret_cast(&previous1_lo[0])), "r"(*reinterpret_cast(&previous1_lo[1])), "r"(*reinterpret_cast(&previous1_lo[2])), "r"(*reinterpret_cast(&previous1_lo[3])), "r"(*reinterpret_cast(&previous1_lo[4])), "r"(*reinterpret_cast(&previous1_lo[5])), "r"(*reinterpret_cast(&previous1_lo[6])), "r"(*reinterpret_cast(&previous1_lo[7]))); + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x2.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(taddr + 112 + 1048576), "r"(*reinterpret_cast(&previous1_hi[0])), "r"(*reinterpret_cast(&previous1_hi[1])), "r"(*reinterpret_cast(&previous1_hi[2])), "r"(*reinterpret_cast(&previous1_hi[3])), "r"(*reinterpret_cast(&previous1_hi[4])), "r"(*reinterpret_cast(&previous1_hi[5])), "r"(*reinterpret_cast(&previous1_hi[6])), "r"(*reinterpret_cast(&previous1_hi[7]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + mbarrier_arrive(o_empty1_addr); + } + } + } + mbarrier_wait(scale_ready0_addr, _phase_scale_ready0_0); + _phase_scale_ready0_0 ^= 1; + float _tmem_load_2[8]; + tmem_ld_x8(&_tmem_load_2[0], taddr + 32); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + mbarrier_arrive(stats_empty0_addr); + mbarrier_wait(scale_ready1_addr, _phase_scale_ready1_0); + _phase_scale_ready1_0 ^= 1; + float _tmem_load_3[8]; + tmem_ld_x8(&_tmem_load_3[0], taddr + 64); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + mbarrier_arrive(stats_empty1_addr); + mbarrier_wait(o_full0_addr, _phase_o_full0_0); + _phase_o_full0_0 ^= 1; + if (num_tiles_c > 1) { + mbarrier_wait(o_full1_addr, _phase_o_full1_0); + _phase_o_full1_0 ^= 1; + } + asm volatile("tcgen05.fence::after_thread_sync;"); + float out0_lo[8]; + float out0_hi[8]; + float out1_lo[8]; + float out1_hi[8]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&out0_lo[0])), "=r"(*reinterpret_cast(&out0_lo[1])), "=r"(*reinterpret_cast(&out0_lo[2])), "=r"(*reinterpret_cast(&out0_lo[3])), "=r"(*reinterpret_cast(&out0_lo[4])), "=r"(*reinterpret_cast(&out0_lo[5])), "=r"(*reinterpret_cast(&out0_lo[6])), "=r"(*reinterpret_cast(&out0_lo[7])) + : "r"(taddr + 96)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&out0_hi[0])), "=r"(*reinterpret_cast(&out0_hi[1])), "=r"(*reinterpret_cast(&out0_hi[2])), "=r"(*reinterpret_cast(&out0_hi[3])), "=r"(*reinterpret_cast(&out0_hi[4])), "=r"(*reinterpret_cast(&out0_hi[5])), "=r"(*reinterpret_cast(&out0_hi[6])), "=r"(*reinterpret_cast(&out0_hi[7])) + : "r"(taddr + 96 + 1048576)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&out1_lo[0])), "=r"(*reinterpret_cast(&out1_lo[1])), "=r"(*reinterpret_cast(&out1_lo[2])), "=r"(*reinterpret_cast(&out1_lo[3])), "=r"(*reinterpret_cast(&out1_lo[4])), "=r"(*reinterpret_cast(&out1_lo[5])), "=r"(*reinterpret_cast(&out1_lo[6])), "=r"(*reinterpret_cast(&out1_lo[7])) + : "r"(taddr + 112)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x2.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(*reinterpret_cast(&out1_hi[0])), "=r"(*reinterpret_cast(&out1_hi[1])), "=r"(*reinterpret_cast(&out1_hi[2])), "=r"(*reinterpret_cast(&out1_hi[3])), "=r"(*reinterpret_cast(&out1_hi[4])), "=r"(*reinterpret_cast(&out1_hi[5])), "=r"(*reinterpret_cast(&out1_hi[6])), "=r"(*reinterpret_cast(&out1_hi[7])) + : "r"(taddr + 112 + 1048576)); + int logical_cta_c = ((batch_c * 32 + y_c) * max_num_ctas_q + cta_q_c) * max_num_ctas_kv + cta_kv_c; + int partial_o_base_c = logical_cta_c * 16 * 128; + int stats_base_c = logical_cta_c * 16; + float final_merged_max_c[4]; + float final_local_sum_c[4]; + float final_scale0_c[4]; + float final_scale1_c[4]; + #pragma unroll + for (int head_local_c_4 = 0; head_local_c_4 < 4; head_local_c_4++) { + float _max_30 = max_noftz(_tmem_load_2[head_local_c_4 + 4], _tmem_load_3[head_local_c_4 + 4]); + float merged_max_c = _max_30; + float _exp2_30 = approx_exp2((_tmem_load_2[head_local_c_4 + 4] - merged_max_c) * bmm1_scale_log2); + float weight0_c = ((_tmem_load_2[head_local_c_4 + 4] == -3.4028235e+38f) ? 0.0f : _exp2_30); + float _exp2_31 = approx_exp2((_tmem_load_3[head_local_c_4 + 4] - merged_max_c) * bmm1_scale_log2); + float weight1_c = ((_tmem_load_3[head_local_c_4 + 4] == -3.4028235e+38f) ? 0.0f : _exp2_31); + final_merged_max_c[head_local_c_4] = merged_max_c; + final_local_sum_c[head_local_c_4] = _tmem_load_2[head_local_c_4] * weight0_c + _tmem_load_3[head_local_c_4] * weight1_c; + final_scale0_c[head_local_c_4] = weight0_c * bmm2_scale / 448.0f; + final_scale1_c[head_local_c_4] = weight1_c * bmm2_scale / 448.0f; + } + #pragma unroll + for (int head_local_c_5 = 0; head_local_c_5 < 4; head_local_c_5++) { + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, final_local_sum_c[head_local_c_5], 16); + final_local_sum_c[head_local_c_5] = final_local_sum_c[head_local_c_5] + _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, final_local_sum_c[head_local_c_5], 8); + final_local_sum_c[head_local_c_5] = final_local_sum_c[head_local_c_5] + _shfl_xor_7; + float _shfl_xor_8 = __shfl_xor_sync(0xFFFFFFFF, final_local_sum_c[head_local_c_5], 4); + final_local_sum_c[head_local_c_5] = final_local_sum_c[head_local_c_5] + _shfl_xor_8; + if (lane < 4) { + corr_stats_reduce[warp_c * 16 + corr_head_base_c + head_local_c_5] = final_local_sum_c[head_local_c_5]; + } + } + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float final_global_sum_c[4]; + #pragma unroll + for (int head_local_c_6 = 0; head_local_c_6 < 4; head_local_c_6++) { + int head_offset_c = corr_head_base_c + head_local_c_6; + final_global_sum_c[head_local_c_6] = corr_stats_reduce[head_offset_c] + corr_stats_reduce[16 + head_offset_c] + corr_stats_reduce[32 + head_offset_c] + corr_stats_reduce[48 + head_offset_c]; + } + float merged_o_c[16]; + #pragma unroll + for (int reg_base_c_2 = 0; reg_base_c_2 < 8; reg_base_c_2 += 4) { + const int scale_idx_c = ((reg_base_c_2 == 0) ? 0 : 2); + float2 _f2_60 = make_float2(final_scale0_c[scale_idx_c], final_scale0_c[scale_idx_c + 1]); + float2 _f2_61 = make_float2(final_scale1_c[scale_idx_c], final_scale1_c[scale_idx_c + 1]); + #pragma unroll + for (int pair_off_c = 0; pair_off_c < 4; pair_off_c += 2) { + const int reg_c = reg_base_c_2 + pair_off_c; + float2 _f2_62 = make_float2(out0_lo[reg_c], out0_lo[reg_c + 1]); + float2 _f2_63 = make_float2(out1_lo[reg_c], out1_lo[reg_c + 1]); + float2 _f2_64 = make_float2(out0_hi[reg_c], out0_hi[reg_c + 1]); + float2 _f2_65 = make_float2(out1_hi[reg_c], out1_hi[reg_c + 1]); + float2 _mul_f32x2_22; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_22) : "l"(*(const unsigned long long*)&_f2_63), "l"(*(const unsigned long long*)&_f2_61)); + float2 out1_lo_scaled_c = _mul_f32x2_22; + float2 _mul_f32x2_23; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_23) : "l"(*(const unsigned long long*)&_f2_65), "l"(*(const unsigned long long*)&_f2_61)); + float2 out1_hi_scaled_c = _mul_f32x2_23; + float2 merged_lo_pair_c = fma_f32x2_rn_ftz(_f2_62, _f2_60, out1_lo_scaled_c); + float2 merged_hi_pair_c = fma_f32x2_rn_ftz(_f2_64, _f2_60, out1_hi_scaled_c); + merged_o_c[reg_c] = merged_lo_pair_c.x; + merged_o_c[reg_c + 1] = merged_lo_pair_c.y; + merged_o_c[reg_c + 8] = merged_hi_pair_c.x; + merged_o_c[reg_c + 9] = merged_hi_pair_c.y; + } + } + unsigned int regs_o_c[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(merged_o_c[_lp*2 + 0], merged_o_c[_lp*2+1 + 0])); + regs_o_c[_lp] = *(uint32_t*)&_bf2; + } + int slice_idx_c = warp_c / 2; + int warp_idx_in_slice_c = warp_c % 2; + int mtx_idx_c = lane / 8; + int mtx_row_idx_c = mtx_idx_c / 2; + int mtx_col_idx_c = warp_idx_in_slice_c * 4 + mtx_idx_c % 2; + int thr_row_idx_c = lane % 8; + #pragma unroll + for (int stsm_idx_c = 0; stsm_idx_c < 2; stsm_idx_c++) { + int seg_col_idx_c = mtx_col_idx_c + stsm_idx_c * 2 ^ thr_row_idx_c; + int stsm_offset_c = slice_idx_c * 16 * 128 + (mtx_row_idx_c * 8 + thr_row_idx_c) * 128 + seg_col_idx_c * 16; + const int stsm_reg_c = stsm_idx_c * 4; + const void* _stmatrix_ptr_8 = reinterpret_cast(reinterpret_cast(smem_o) + stsm_offset_c); + uint64_t _stmatrix_addr64_8; + asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(_stmatrix_addr64_8) : "l"(_stmatrix_ptr_8)); + uint32_t _stmatrix_addr_8; + asm volatile("cvt.u32.u64 %0, %1;" : "=r"(_stmatrix_addr_8) : "l"(_stmatrix_addr64_8)); + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_8), "r"(*reinterpret_cast(®s_o_c[stsm_reg_c])), "r"(*reinterpret_cast(®s_o_c[stsm_reg_c + 1])), "r"(*reinterpret_cast(®s_o_c[stsm_reg_c + 2])), "r"(*reinterpret_cast(®s_o_c[stsm_reg_c + 3])) + : "memory"); + } + asm volatile("barrier.sync 3, 128;" ::: "memory"); + #pragma unroll + for (int copy_idx_c = 0; copy_idx_c < 2; copy_idx_c++) { + int copy_base_c = corr_tid * 16 + copy_idx_c * 128 * 16; + int copy_smem_row_c = copy_base_c / 128; + int copy_dst_row_c = copy_smem_row_c % 16; + int copy_dst_col_bytes_c = copy_smem_row_c / 16 * 128 + copy_base_c % 128; + int copy_smem_offset_c = copy_base_c ^ copy_smem_row_c % 8 * 16; + unsigned int copy_vec_c[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(©_vec_c[0])), "=r"(*reinterpret_cast(©_vec_c[(0) + 1])), "=r"(*reinterpret_cast(©_vec_c[(0) + 2])), "=r"(*reinterpret_cast(©_vec_c[(0) + 3])) + : "r"(smem_o_addr + (unsigned int)copy_smem_offset_c)); + reinterpret_cast(partial_O + (partial_o_base_c + copy_dst_row_c * 128 + copy_dst_col_bytes_c / 2))[0] = reinterpret_cast(copy_vec_c)[0]; + } + if (corr_tid < 4) { + #pragma unroll + for (int head_local_c_7 = 0; head_local_c_7 < 4; head_local_c_7++) { + int head_c = 2 * (corr_tid % 4) + 8 * (head_local_c_7 / 2) + head_local_c_7 % 2; + partial_stats[(stats_base_c + head_c) * 2] = final_merged_max_c[head_local_c_7]; + partial_stats[(stats_base_c + head_c) * 2 + 1] = final_global_sum_c[head_local_c_7]; + } + } + mbarrier_arrive(o_empty0_addr); + if (num_tiles_c > 1) { + mbarrier_arrive(o_empty1_addr); + } + int reduction_group_c = (batch_c * 32 + y_c) * max_num_ctas_q + cta_q_c; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + if (corr_tid == 0) { + uint32_t _atomic_inc_old_0; + asm volatile("atom.acq_rel.gpu.global.inc.u32 %0, [%1], %2;" + : "=r"(_atomic_inc_old_0) : "l"(&completion[reduction_group_c]), "r"(static_cast(1)) : "memory"); + unsigned int old_ticket_c = _atomic_inc_old_0; + reduction_ticket[0] = 1 - (int)old_ticket_c; + if (old_ticket_c == 0) { + { + unsigned int* _global_wait_zero_p_9 = (reinterpret_cast(completion) + (reduction_group_c)); + unsigned int _global_wait_zero_v_9; + do { + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(_global_wait_zero_v_9) : "l"(_global_wait_zero_p_9) : "memory"); + } while (_global_wait_zero_v_9 != 0u); + } + } + } + asm volatile("barrier.sync 3, 128;" ::: "memory"); + int reduction_cta_c = reduction_ticket[0]; + int reduction_row_c = reduction_cta_c * 8 + corr_tid / 16; + int reduction_col_c = corr_tid * 8 % 128; + int stats_group_base_c = reduction_group_c * max_num_ctas_kv * 16 * 2; + int stats_row_c = reduction_row_c * 2; + float _vec_load_0[2]; + { + float2 _v2_10 = *reinterpret_cast(partial_stats + stats_group_base_c + stats_row_c); + _vec_load_0[0] = _v2_10.x; + _vec_load_0[0 + 1] = _v2_10.y; + } + float _vec_load_1[2]; + { + float2 _v2_11 = *reinterpret_cast(partial_stats + stats_group_base_c + 32 + stats_row_c); + _vec_load_1[0] = _v2_11.x; + _vec_load_1[0 + 1] = _v2_11.y; + } + float max0_c = _vec_load_0[0]; + float sum0_c = _vec_load_0[1]; + float max1_c = _vec_load_1[0]; + float sum1_c = _vec_load_1[1]; + float _max_31 = max_noftz(max0_c, max1_c); + float global_max_c = _max_31; + float _exp2_32 = approx_exp2((max0_c - global_max_c) * bmm1_scale_log2); + float cta_scale0_c = _exp2_32; + float _exp2_33 = approx_exp2((max1_c - global_max_c) * bmm1_scale_log2); + float cta_scale1_c = _exp2_33; + float global_sum_c = sum0_c * cta_scale0_c + sum1_c * cta_scale1_c; + float normalized_scale_c = 448.0f / global_sum_c; + int partial_group_base_c = reduction_group_c * max_num_ctas_kv * 16 * 128; + int partial_row_col_c = reduction_row_c * 128 + reduction_col_c; + float _vec_load_2[8]; + { + const uint4* _vptr_12 = reinterpret_cast(partial_O + (partial_group_base_c + partial_row_col_c) + 0); + uint4 _vld_12[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_12[_blk] = _vptr_12[_blk]; + uint32_t* _vpairs_12 = reinterpret_cast(&_vld_12[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&_vec_load_2[0 + _blk * 8 + _pair * 2])[0]), "=f"((&_vec_load_2[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_12[_pair])); + } + } + } + float _vec_load_3[8]; + { + const uint4* _vptr_13 = reinterpret_cast(partial_O + (partial_group_base_c + 2048 + partial_row_col_c) + 0); + uint4 _vld_13[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_13[_blk] = _vptr_13[_blk]; + uint32_t* _vpairs_13 = reinterpret_cast(&_vld_13[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&_vec_load_3[0 + _blk * 8 + _pair * 2])[0]), "=f"((&_vec_load_3[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_13[_pair])); + } + } + } + float reduced_c[8]; + float2 _f2_66 = make_float2(cta_scale0_c, cta_scale0_c); + float2 _f2_67 = make_float2(cta_scale1_c, cta_scale1_c); + float2 _f2_68 = make_float2(normalized_scale_c, normalized_scale_c); + #pragma unroll + for (int elem_c = 0; elem_c < 8; elem_c += 2) { + float2 _f2_69 = make_float2(_vec_load_2[elem_c], _vec_load_2[elem_c + 1]); + float2 _f2_70 = make_float2(_vec_load_3[elem_c], _vec_load_3[elem_c + 1]); + float2 _mul_f32x2_24; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_24) : "l"(*(const unsigned long long*)&_f2_70), "l"(*(const unsigned long long*)&_f2_67)); + float2 merged_pair_c = fma_f32x2_rn_ftz(_f2_69, _f2_66, _mul_f32x2_24); + float2 _mul_f32x2_25; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_25) : "l"(*(const unsigned long long*)&merged_pair_c), "l"(*(const unsigned long long*)&_f2_68)); + float2 reduced_pair_c = _mul_f32x2_25; + reduced_c[elem_c] = reduced_pair_c.x; + reduced_c[elem_c + 1] = reduced_pair_c.y; + } + int output_head_c = head_group_c * 16 + reduction_row_c; + int output_base_c = ((batch_c * q_len + cta_q_c) * 128 + output_head_c) * 512 + hv_partition_c * 128 + reduction_col_c; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(reduced_c[0 + 0], reduced_c[0 + 1]); + _pk[1] = __floats2bfloat162_rn(reduced_c[0 + 2], reduced_c[0 + 3]); + _pk[2] = __floats2bfloat162_rn(reduced_c[0 + 4], reduced_c[0 + 5]); + _pk[3] = __floats2bfloat162_rn(reduced_c[0 + 6], reduced_c[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + output_base_c))[0]) = *reinterpret_cast(&_pk[0]); + } + asm volatile("barrier.sync 13, 128;" ::: "memory"); + if (warp_c == 0) { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: mma_warp ---- + } else if (warp == 12) { + { // mma_warp_main + int cta_q_m = blockIdx.x / 2; + int batch_m = blockIdx.z; + int seq_len_m = seq_lens[batch_m] - (q_len - 1 - cta_q_m); + int num_ctas_m = 2; + int num_steps_m = 2; + int kv_stage_m = 0; + unsigned int _phase_q_full_0 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + unsigned int _phase_s_empty0_0 = 1; + mbarrier_wait(s_empty0_addr, _phase_s_empty0_0); + _phase_s_empty0_0 ^= 1; + unsigned int _phase_kv_full = 0; + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_0 = make_warp_uniform((((smem_q0_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134479888, 0); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_1 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_1 = make_warp_uniform((((smem_q1_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_2 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_2 = make_warp_uniform((((smem_q2_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_2); + uint64_t _mma_ss_b_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_2); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_3 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_3 = make_warp_uniform((((smem_q3_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_3); + uint64_t _mma_ss_b_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_3); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_4 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_4 = make_warp_uniform((((smem_q4_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_4); + uint64_t _mma_ss_b_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_4); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + elect_commit(s_full0_addr); + unsigned int _phase_s_empty1_0 = 1; + mbarrier_wait(s_empty1_addr, _phase_s_empty1_0); + _phase_s_empty1_0 ^= 1; + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_5 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_5 = make_warp_uniform((((smem_q0_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_5); + uint64_t _mma_ss_b_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_5); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134479888, 0); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_6 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_6 = make_warp_uniform((((smem_q1_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_6 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_6); + uint64_t _mma_ss_b_desc_6 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_6); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_6, _mma_ss_b_desc_6, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_6, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_6, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_6, _mma_ss_b_desc_6, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_6, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_6, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_6, _mma_ss_b_desc_6, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_6, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_6, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_6, _mma_ss_b_desc_6, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_7 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_7 = make_warp_uniform((((smem_q2_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_7 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_7); + uint64_t _mma_ss_b_desc_7 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_7); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_7, _mma_ss_b_desc_7, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_7, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_7, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_7, _mma_ss_b_desc_7, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_7, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_7, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_7, _mma_ss_b_desc_7, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_7, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_7, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_7, _mma_ss_b_desc_7, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_8 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_8 = make_warp_uniform((((smem_q3_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_8 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_8); + uint64_t _mma_ss_b_desc_8 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_8); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_8, _mma_ss_b_desc_8, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_8, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_8, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_8, _mma_ss_b_desc_8, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_8, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_8, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_8, _mma_ss_b_desc_8, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_8, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_8, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_8, _mma_ss_b_desc_8, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_9 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_9 = make_warp_uniform((((smem_q4_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_9 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_9); + uint64_t _mma_ss_b_desc_9 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_9); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_9, _mma_ss_b_desc_9, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_9, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_9, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_9, _mma_ss_b_desc_9, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_9, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_9, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_9, _mma_ss_b_desc_9, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_9, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_9, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_9, _mma_ss_b_desc_9, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + elect_commit(s_full1_addr); + unsigned int _phase_o_empty0_0 = 1; + unsigned int _phase_o_empty1_0 = 1; + #pragma unroll 1 + for (int pair_m = 0; pair_m < num_steps_m - 1; pair_m++) { + mbarrier_wait(s_empty0_addr, _phase_s_empty0_0); + _phase_s_empty0_0 ^= 1; + mbarrier_wait(o_empty0_addr, _phase_o_empty0_0); + _phase_o_empty0_0 ^= 1; + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_10 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (kv_stage_m) * 1024); + int _mma_b_lo_10 = make_warp_uniform((((smem_p0_addr) >> 4) & 0x3FFF) | 0x800000); + { + uint64_t _mma_ss_a_desc_10 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_10); + uint64_t _mma_ss_b_desc_10 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_10); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_10, _mma_ss_b_desc_10, 134512656, ((pair_m == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_10, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_10, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_10, _mma_ss_b_desc_10, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_10, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_10, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_10, _mma_ss_b_desc_10, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_10, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_10, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_10, _mma_ss_b_desc_10, 134512656, 1); + } + } + elect_commit2(kv_empty_addr + (kv_stage_m) * 8, o_full0_addr); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_11 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_11 = make_warp_uniform((((smem_q0_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_11 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_11); + uint64_t _mma_ss_b_desc_11 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_11); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_11, _mma_ss_b_desc_11, 134479888, 0); + } + incr_smem_desc_lo(_mma_ss_a_desc_11, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_11, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_11, _mma_ss_b_desc_11, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_11, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_11, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_11, _mma_ss_b_desc_11, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_11, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_11, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_11, _mma_ss_b_desc_11, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_12 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_12 = make_warp_uniform((((smem_q1_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_12 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_12); + uint64_t _mma_ss_b_desc_12 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_12); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_12, _mma_ss_b_desc_12, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_12, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_12, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_12, _mma_ss_b_desc_12, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_12, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_12, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_12, _mma_ss_b_desc_12, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_12, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_12, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_12, _mma_ss_b_desc_12, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_13 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_13 = make_warp_uniform((((smem_q2_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_13 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_13); + uint64_t _mma_ss_b_desc_13 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_13); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_13, _mma_ss_b_desc_13, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_13, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_13, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_13, _mma_ss_b_desc_13, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_13, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_13, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_13, _mma_ss_b_desc_13, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_13, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_13, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_13, _mma_ss_b_desc_13, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_14 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_14 = make_warp_uniform((((smem_q3_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_14 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_14); + uint64_t _mma_ss_b_desc_14 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_14); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_14, _mma_ss_b_desc_14, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_14, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_14, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_14, _mma_ss_b_desc_14, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_14, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_14, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_14, _mma_ss_b_desc_14, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_14, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_14, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_14, _mma_ss_b_desc_14, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_15 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_15 = make_warp_uniform((((smem_q4_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_15 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_15); + uint64_t _mma_ss_b_desc_15 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_15); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_15, _mma_ss_b_desc_15, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_15, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_15, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_15, _mma_ss_b_desc_15, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_15, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_15, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_15, _mma_ss_b_desc_15, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_15, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_15, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4(tmem_tmem, _mma_ss_a_desc_15, _mma_ss_b_desc_15, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + elect_commit(s_full0_addr); + mbarrier_wait(s_empty1_addr, _phase_s_empty1_0); + _phase_s_empty1_0 ^= 1; + mbarrier_wait(o_empty1_addr, _phase_o_empty1_0); + _phase_o_empty1_0 ^= 1; + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_16 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (kv_stage_m) * 1024); + int _mma_b_lo_16 = make_warp_uniform((((smem_p1_addr) >> 4) & 0x3FFF) | 0x800000); + { + uint64_t _mma_ss_a_desc_16 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_16); + uint64_t _mma_ss_b_desc_16 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_16); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_16, _mma_ss_b_desc_16, 134512656, ((pair_m == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_16, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_16, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_16, _mma_ss_b_desc_16, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_16, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_16, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_16, _mma_ss_b_desc_16, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_16, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_16, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_16, _mma_ss_b_desc_16, 134512656, 1); + } + } + elect_commit2(kv_empty_addr + (kv_stage_m) * 8, o_full1_addr); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_17 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_17 = make_warp_uniform((((smem_q0_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_17 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_17); + uint64_t _mma_ss_b_desc_17 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_17); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_17, _mma_ss_b_desc_17, 134479888, 0); + } + incr_smem_desc_lo(_mma_ss_a_desc_17, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_17, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_17, _mma_ss_b_desc_17, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_17, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_17, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_17, _mma_ss_b_desc_17, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_17, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_17, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_17, _mma_ss_b_desc_17, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_18 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_18 = make_warp_uniform((((smem_q1_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_18 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_18); + uint64_t _mma_ss_b_desc_18 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_18); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_18, _mma_ss_b_desc_18, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_18, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_18, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_18, _mma_ss_b_desc_18, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_18, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_18, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_18, _mma_ss_b_desc_18, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_18, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_18, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_18, _mma_ss_b_desc_18, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_19 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_19 = make_warp_uniform((((smem_q2_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_19 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_19); + uint64_t _mma_ss_b_desc_19 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_19); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_19, _mma_ss_b_desc_19, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_19, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_19, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_19, _mma_ss_b_desc_19, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_19, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_19, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_19, _mma_ss_b_desc_19, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_19, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_19, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_19, _mma_ss_b_desc_19, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_20 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_20 = make_warp_uniform((((smem_q3_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_20 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_20); + uint64_t _mma_ss_b_desc_20 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_20); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_20, _mma_ss_b_desc_20, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_20, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_20, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_20, _mma_ss_b_desc_20, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_20, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_20, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_20, _mma_ss_b_desc_20, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_20, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_20, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_20, _mma_ss_b_desc_20, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_21 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage_m) * 1024); + int _mma_b_lo_21 = make_warp_uniform((((smem_q4_addr) >> 4) & 0x3FFF) + (0) * 640); + { + uint64_t _mma_ss_a_desc_21 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_21); + uint64_t _mma_ss_b_desc_21 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_21); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_21, _mma_ss_b_desc_21, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_21, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_21, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_21, _mma_ss_b_desc_21, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_21, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_21, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_21, _mma_ss_b_desc_21, 134479888, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_21, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_21, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (16)), _mma_ss_a_desc_21, _mma_ss_b_desc_21, 134479888, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage_m) * 8); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + elect_commit(s_full1_addr); + } + mbarrier_wait(s_empty0_addr, _phase_s_empty0_0); + _phase_s_empty0_0 ^= 1; + mbarrier_wait(o_empty0_addr, _phase_o_empty0_0); + _phase_o_empty0_0 ^= 1; + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_22 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (kv_stage_m) * 1024); + int _mma_b_lo_22 = make_warp_uniform((((smem_p0_addr) >> 4) & 0x3FFF) | 0x800000); + { + uint64_t _mma_ss_a_desc_22 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_22); + uint64_t _mma_ss_b_desc_22 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_22); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_22, _mma_ss_b_desc_22, 134512656, ((num_steps_m == 1) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_22, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_22, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_22, _mma_ss_b_desc_22, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_22, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_22, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_22, _mma_ss_b_desc_22, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_22, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_22, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (96)), _mma_ss_a_desc_22, _mma_ss_b_desc_22, 134512656, 1); + } + } + elect_commit2(kv_empty_addr + (kv_stage_m) * 8, o_full0_addr); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + mbarrier_wait(s_empty1_addr, _phase_s_empty1_0); + _phase_s_empty1_0 ^= 1; + mbarrier_wait(o_empty1_addr, _phase_o_empty1_0); + _phase_o_empty1_0 ^= 1; + mbarrier_wait(kv_full_addr + (kv_stage_m) * 8, _phase_kv_full); + int _mma_a_lo_23 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (kv_stage_m) * 1024); + int _mma_b_lo_23 = make_warp_uniform((((smem_p1_addr) >> 4) & 0x3FFF) | 0x800000); + { + uint64_t _mma_ss_a_desc_23 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_23); + uint64_t _mma_ss_b_desc_23 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_23); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_23, _mma_ss_b_desc_23, 134512656, ((num_steps_m == 1) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_23, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_23, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_23, _mma_ss_b_desc_23, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_23, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_23, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_23, _mma_ss_b_desc_23, 134512656, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_23, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_23, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (112)), _mma_ss_a_desc_23, _mma_ss_b_desc_23, 134512656, 1); + } + } + elect_commit2(kv_empty_addr + (kv_stage_m) * 8, o_full1_addr); + kv_stage_m += 1; + if (kv_stage_m == 9) { kv_stage_m = 0; _phase_kv_full ^= 1; } + elect_commit(q_empty_addr); + } + // ---- Role: page_warp ---- + } else if (warp == 13) { + { // page_warp_main + int cta_q_p = blockIdx.x / 2; + int cta_kv_p = blockIdx.x % 2; + int batch_p = blockIdx.z; + int seq_len_p = seq_lens[batch_p] - (q_len - 1 - cta_q_p); + int num_ctas_p = 2; + int num_steps_p = 2; + int num_tiles_p = num_steps_p * 2; + int max_page_p = (seq_len_p + 32 - 1) / 32 - 1; + int page_stage_p = 0; + unsigned int _phase_page_empty = 1; + #pragma unroll 1 + for (int tile_p = 0; tile_p < num_tiles_p; tile_p++) { + mbarrier_wait(page_empty_addr + (page_stage_p) * 8, _phase_page_empty); + int global_page_p = (unsigned int)((cta_kv_p * num_tiles_p + tile_p) * 4) + lane; + if (global_page_p > max_page_p) { + global_page_p = max_page_p; + } + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 4;" + :: "r"(page_offsets_addr + (unsigned int)(page_stage_p * 128) + lane * 4), "l"(page_table + (batch_p * page_table_stride + global_page_p))); + asm volatile( + "{\n\t" + "cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n\t" + "}" + :: "r"(page_full_addr + (page_stage_p) * 8) : "memory"); + page_stage_p += 1; + if (page_stage_p == 6) { page_stage_p = 0; _phase_page_empty ^= 1; } + } + } + // ---- Role: padding_warp ---- + } else if (warp == 14) { + // idle — no tasks assigned + // ---- Role: load_warp ---- + } else if (warp == 15) { + { // load_warp_main + int cta_q_l = blockIdx.x / 2; + int cta_kv_l = blockIdx.x % 2; + int y_l = blockIdx.y; + int head_group_l = y_l / 4; + int hv_partition_l = y_l % 4; + int batch_l = blockIdx.z; + int seq_len_l = seq_lens[batch_l] - (q_len - 1 - cta_q_l); + int num_ctas_l = 2; + int num_steps_l = 2; + int num_tiles_l = num_steps_l * 2; + int q_stage_l = 0; + int kv_stage_l = 0; + int page_stage_l = 0; + unsigned int _phase_q_empty = 1; + mbarrier_wait(q_empty_addr + (q_stage_l) * 8, _phase_q_empty); + if (elect_sync()) { + int q_row_l = (batch_l * q_len + cta_q_l) * 128 + head_group_l * 16; + mbarrier_arrive_expect_tx(q_full_addr + (q_stage_l) * 8, 10240); + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;" + :: "r"(smem_q0_addr + (unsigned int)(q_stage_l * 10240)), "l"((&Q_map)), "r"(0), "r"(q_row_l), + "r"(q_full_addr + (q_stage_l) * 8), "l"(0x14F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;" + :: "r"(smem_q1_addr + (unsigned int)(q_stage_l * 10240)), "l"((&Q_map)), "r"(128), "r"(q_row_l), + "r"(q_full_addr + (q_stage_l) * 8), "l"(0x14F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;" + :: "r"(smem_q2_addr + (unsigned int)(q_stage_l * 10240)), "l"((&Q_map)), "r"(256), "r"(q_row_l), + "r"(q_full_addr + (q_stage_l) * 8), "l"(0x14F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;" + :: "r"(smem_q3_addr + (unsigned int)(q_stage_l * 10240)), "l"((&Q_map)), "r"(384), "r"(q_row_l), + "r"(q_full_addr + (q_stage_l) * 8), "l"(0x14F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;" + :: "r"(smem_q4_addr + (unsigned int)(q_stage_l * 10240)), "l"((&Q_map)), "r"(512), "r"(q_row_l), + "r"(q_full_addr + (q_stage_l) * 8), "l"(0x14F0000000000000ULL) : "memory"); + } + unsigned int _phase_kv_empty = 1; + #pragma unroll 1 + for (int aligned_tile_l = 0; aligned_tile_l < num_tiles_l; aligned_tile_l++) { + mbarrier_wait(page_full_addr + (aligned_tile_l) * 8, 0); + int cached_page_indices[4]; + if (elect_sync()) { + #pragma unroll + for (int cached_page = 0; cached_page < 4; cached_page++) { + cached_page_indices[cached_page] = page_offsets[aligned_tile_l * 32 + cached_page]; + } + } + #pragma unroll + for (int head_stage = 0; head_stage < 5; head_stage++) { + mbarrier_wait(kv_empty_addr + (kv_stage_l) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_l) * 8, 16384); + #pragma unroll + for (int page = 0; page < 4; page++) { + int cached_physical_page = cached_page_indices[page]; + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;" + :: "r"(smem_kv_addr + (unsigned int)(kv_stage_l * 16384) + (unsigned int)(page * 32 * 128)), "l"((&K_map)), "r"(head_stage * 128), "r"(cached_physical_page * 32), + "r"(kv_full_addr + (kv_stage_l) * 8), "l"(0x14F0000000000000ULL) : "memory"); + } + } + kv_stage_l += 1; + if (kv_stage_l == 9) { kv_stage_l = 0; _phase_kv_empty ^= 1; } + } + if (aligned_tile_l > 0) { + mbarrier_wait(kv_empty_addr + (kv_stage_l) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_l) * 8, 16384); + #pragma unroll + for (int page_1 = 0; page_1 < 4; page_1++) { + int physical_page = page_offsets[(aligned_tile_l - 1) * 32 + page_1]; + tma_2d_gmem2smem(smem_kv_addr + (unsigned int)(kv_stage_l * 16384) + (unsigned int)(page_1 * 32 * 128), (&K_map), hv_partition_l * 128, physical_page * 32, kv_full_addr + (kv_stage_l) * 8); + } + } + kv_stage_l += 1; + if (kv_stage_l == 9) { kv_stage_l = 0; _phase_kv_empty ^= 1; } + mbarrier_arrive(page_empty_addr + (aligned_tile_l - 1) * 8); + } + } + mbarrier_wait(kv_empty_addr + (kv_stage_l) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_l) * 8, 16384); + #pragma unroll + for (int page_2 = 0; page_2 < 4; page_2++) { + int physical_page_1 = page_offsets[(num_tiles_l - 1) * 32 + page_2]; + tma_2d_gmem2smem(smem_kv_addr + (unsigned int)(kv_stage_l * 16384) + (unsigned int)(page_2 * 32 * 128), (&K_map), hv_partition_l * 128, physical_page_1 * 32, kv_full_addr + (kv_stage_l) * 8); + } + } + kv_stage_l += 1; + if (kv_stage_l == 9) { kv_stage_l = 0; _phase_kv_empty ^= 1; } + mbarrier_arrive(page_empty_addr + (num_tiles_l - 1) * 8); + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_full_abi_tail_bf16_83451b7317.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_full_abi_tail_bf16_83451b7317.cu new file mode 100644 index 00000000000..b8c039f2051 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_full_abi_tail_bf16_83451b7317.cu @@ -0,0 +1,195 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_WEIGHTS_OFF 0 +#define SMEM_WEIGHTS_STAGE_BYTES 16384 +#define SMEM_WEIGHTS_STRIDE 16384 +#define SMEM_TOTAL 16384 +#define THREADS 32 + +#include + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + +extern "C" { + +__global__ __launch_bounds__(32, 1) void +kernel_full_abi_tail_bf16(__nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ KV, float* __restrict__ fp8_lut, int* __restrict__ page_table, int* __restrict__ sparse_indices, int* __restrict__ row_batches, int* __restrict__ row_seq_lens, __nv_bfloat16* __restrict__ O, float* __restrict__ LSE, float* __restrict__ sinks, int num_heads, int qk_dim, int value_dim, int kv_stride, int page_size, int page_table_width, int sparse_width, int use_sparse, float softmax_scale, float bmm2_scale, int enable_sink, int write_lse) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + float* weights = reinterpret_cast(smem_raw + 0); + const int weights_addr = smem + 0; + + // === Task calls (dependency order) === + int row = blockIdx.x; + int head = blockIdx.y; + int batch = row_batches[row]; + int seq_len = row_seq_lens[row]; + int q_base = (row * num_heads + head) * qk_dim; + float local_max = -MLA_INF; + #pragma unroll 1 + for (int token = lane; token < seq_len; token += 32) { + int physical_token = 0; + if (use_sparse != 0) { + physical_token = sparse_indices[row * sparse_width + token]; + } else { + int logical_page = token / page_size; + int token_in_page = token - logical_page * page_size; + int physical_page = page_table[batch * page_table_width + logical_page]; + physical_token = physical_page * page_size + token_in_page; + } + int kv_base = physical_token * kv_stride; + float score = 0.0f; + #pragma unroll 4 + for (int dim = 0; dim < qk_dim; dim++) { + float q_value = Q[q_base + dim]; + float kv_value = KV[kv_base + dim]; + score = score + q_value * kv_value; + } + score = score * softmax_scale; + weights[token] = score; + float _max_0 = max_noftz(local_max, score); + local_max = _max_0; + } + if (enable_sink != 0 && lane == 0) { + float _max_1 = max_noftz(local_max, sinks[head]); + local_max = _max_1; + } + float _warp_reduce_0 = local_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + float row_max = _warp_reduce_0; + __syncthreads(); + float local_sum = 0.0f; + #pragma unroll 1 + for (int token_1 = lane; token_1 < seq_len; token_1 += 32) { + float _exp2_0 = approx_exp2((weights[token_1] - row_max) * 1.4426950408889634f); + float weight = _exp2_0; + weights[token_1] = weight; + local_sum = local_sum + weight; + } + if (enable_sink != 0 && lane == 0) { + float _exp2_1 = approx_exp2((sinks[head] - row_max) * 1.4426950408889634f); + local_sum = local_sum + _exp2_1; + } + float _warp_reduce_1 = local_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + float row_sum = _warp_reduce_1; + float _rcp_0 = approx_rcp(row_sum); + float inv_sum = _rcp_0; + __syncthreads(); + int o_base = (row * num_heads + head) * value_dim; + #pragma unroll 1 + for (int dim_1 = lane; dim_1 < value_dim; dim_1 += 32) { + float accum = 0.0f; + #pragma unroll 1 + for (int token_2 = 0; token_2 < seq_len; token_2++) { + int physical_token_1 = 0; + if (use_sparse != 0) { + physical_token_1 = sparse_indices[row * sparse_width + token_2]; + } else { + int logical_page_1 = token_2 / page_size; + int token_in_page_1 = token_2 - logical_page_1 * page_size; + int physical_page_1 = page_table[batch * page_table_width + logical_page_1]; + physical_token_1 = physical_page_1 * page_size + token_in_page_1; + } + int kv_base_1 = physical_token_1 * kv_stride; + float value = KV[kv_base_1 + dim_1]; + accum = accum + weights[token_2] * value; + } + *(reinterpret_cast<__nv_bfloat16*>(O + (o_base + dim_1)) + (0)) = __float2bfloat16_rn(accum * inv_sum * bmm2_scale); + } + if (write_lse != 0 && lane == 0) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(row_sum)); + float lse = row_max * 1.4426950408889634f + _log2_0; + *(reinterpret_cast(LSE + (row * num_heads + head)) + (0)) = lse; + } +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_full_abi_tail_fp8_26b36022eb.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_full_abi_tail_fp8_26b36022eb.cu new file mode 100644 index 00000000000..58e77e5a500 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_full_abi_tail_fp8_26b36022eb.cu @@ -0,0 +1,195 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_WEIGHTS_OFF 0 +#define SMEM_WEIGHTS_STAGE_BYTES 16384 +#define SMEM_WEIGHTS_STRIDE 16384 +#define SMEM_TOTAL 16384 +#define THREADS 32 + +#include + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + +extern "C" { + +__global__ __launch_bounds__(32, 1) void +kernel_full_abi_tail_fp8(uint8_t* __restrict__ Q, uint8_t* __restrict__ KV, float* __restrict__ fp8_lut, int* __restrict__ page_table, int* __restrict__ sparse_indices, int* __restrict__ row_batches, int* __restrict__ row_seq_lens, __nv_bfloat16* __restrict__ O, float* __restrict__ LSE, float* __restrict__ sinks, int num_heads, int qk_dim, int value_dim, int kv_stride, int page_size, int page_table_width, int sparse_width, int use_sparse, float softmax_scale, float bmm2_scale, int enable_sink, int write_lse) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + float* weights = reinterpret_cast(smem_raw + 0); + const int weights_addr = smem + 0; + + // === Task calls (dependency order) === + int row = blockIdx.x; + int head = blockIdx.y; + int batch = row_batches[row]; + int seq_len = row_seq_lens[row]; + int q_base = (row * num_heads + head) * qk_dim; + float local_max = -MLA_INF; + #pragma unroll 1 + for (int token = lane; token < seq_len; token += 32) { + int physical_token = 0; + if (use_sparse != 0) { + physical_token = sparse_indices[row * sparse_width + token]; + } else { + int logical_page = token / page_size; + int token_in_page = token - logical_page * page_size; + int physical_page = page_table[batch * page_table_width + logical_page]; + physical_token = physical_page * page_size + token_in_page; + } + int kv_base = physical_token * kv_stride; + float score = 0.0f; + #pragma unroll 4 + for (int dim = 0; dim < qk_dim; dim++) { + float q_value = fp8_lut[Q[q_base + dim]]; + float kv_value = fp8_lut[KV[kv_base + dim]]; + score = score + q_value * kv_value; + } + score = score * softmax_scale; + weights[token] = score; + float _max_0 = max_noftz(local_max, score); + local_max = _max_0; + } + if (enable_sink != 0 && lane == 0) { + float _max_1 = max_noftz(local_max, sinks[head]); + local_max = _max_1; + } + float _warp_reduce_0 = local_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + float row_max = _warp_reduce_0; + __syncthreads(); + float local_sum = 0.0f; + #pragma unroll 1 + for (int token_1 = lane; token_1 < seq_len; token_1 += 32) { + float _exp2_0 = approx_exp2((weights[token_1] - row_max) * 1.4426950408889634f); + float weight = _exp2_0; + weights[token_1] = weight; + local_sum = local_sum + weight; + } + if (enable_sink != 0 && lane == 0) { + float _exp2_1 = approx_exp2((sinks[head] - row_max) * 1.4426950408889634f); + local_sum = local_sum + _exp2_1; + } + float _warp_reduce_1 = local_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + float row_sum = _warp_reduce_1; + float _rcp_0 = approx_rcp(row_sum); + float inv_sum = _rcp_0; + __syncthreads(); + int o_base = (row * num_heads + head) * value_dim; + #pragma unroll 1 + for (int dim_1 = lane; dim_1 < value_dim; dim_1 += 32) { + float accum = 0.0f; + #pragma unroll 1 + for (int token_2 = 0; token_2 < seq_len; token_2++) { + int physical_token_1 = 0; + if (use_sparse != 0) { + physical_token_1 = sparse_indices[row * sparse_width + token_2]; + } else { + int logical_page_1 = token_2 / page_size; + int token_in_page_1 = token_2 - logical_page_1 * page_size; + int physical_page_1 = page_table[batch * page_table_width + logical_page_1]; + physical_token_1 = physical_page_1 * page_size + token_in_page_1; + } + int kv_base_1 = physical_token_1 * kv_stride; + float value = fp8_lut[KV[kv_base_1 + dim_1]]; + accum = accum + weights[token_2] * value; + } + *(reinterpret_cast<__nv_bfloat16*>(O + (o_base + dim_1)) + (0)) = __float2bfloat16_rn(accum * inv_sum * bmm2_scale); + } + if (write_lse != 0 && lane == 0) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(row_sum)); + float lse = row_max * 1.4426950408889634f + _log2_0; + *(reinterpret_cast(LSE + (row * num_heads + head)) + (0)) = lse; + } +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_5f26657c8f.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_5f26657c8f.cu new file mode 100644 index 00000000000..8979737b965 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_5f26657c8f.cu @@ -0,0 +1,2499 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 +#define EXACT007826_CONSTANTS 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x8(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x8.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * ((EXACT007826_CONSTANTS) ? 16 : value_split_count) + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < ((EXACT007826_CONSTANTS) ? 8192 : total_work_items); work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) > 76) { + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[((EXACT007826_CONSTANTS) ? 8192 : total_work_items) * ((EXACT007826_CONSTANTS) ? 32 : max_pages_per_seq) + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * ((EXACT007826_CONSTANTS) ? 32 : max_pages_per_seq); + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)(((EXACT007826_CONSTANTS) ? 16 : value_split_count)) + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * ((EXACT007826_CONSTANTS) ? 16 : value_split_count) + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < ((EXACT007826_CONSTANTS) ? 8192 : total_work_items); _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)(((EXACT007826_CONSTANTS) ? 16 : value_split_count)) + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 176;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * ((EXACT007826_CONSTANTS) ? 16 : value_split_count) + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < ((EXACT007826_CONSTANTS) ? 8192 : total_work_items); work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) > 76) { + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[((EXACT007826_CONSTANTS) ? 8192 : total_work_items) * ((EXACT007826_CONSTANTS) ? 32 : max_pages_per_seq) + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + uint32_t _mbar_token_10 = mbarrier_try_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + int abs_tile = start_tile_1 + tile_1; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + bool score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + float sv[64]; + float tile_max = -MLA_INF; + float new_max = -MLA_INF; + if (!score_tile_full) { + mbarrier_wait_token(s_full_addr + (phase) * 8, s_wait_phase, _mbar_token_10); + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]), "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]), "=f"(tile_max) + : "r"(s_base)); + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + } else { + mbarrier_wait_token(s_full_addr + (phase) * 8, s_wait_phase, _mbar_token_10); + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]), "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]), "=f"(tile_max) + : "r"(s_base)); + new_max = tile_max; + } + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)(((EXACT007826_CONSTANTS) ? 16 : value_split_count)) + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * ((EXACT007826_CONSTANTS) ? 16 : value_split_count) + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < ((EXACT007826_CONSTANTS) ? 8192 : total_work_items); work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) > 76) { + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[((EXACT007826_CONSTANTS) ? 8192 : total_work_items) * ((EXACT007826_CONSTANTS) ? 32 : max_pages_per_seq) + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll 1 + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 64) { + float _tmem_load_6[32]; + tmem_ld_x8(&_tmem_load_6[0], o_base_epi + c_1); + tmem_ld_x8(&_tmem_load_6[8], o_base_epi + c_1 + 8); + tmem_ld_x8(&_tmem_load_6[16], o_base_epi + c_1 + 16); + tmem_ld_x8(&_tmem_load_6[24], o_base_epi + c_1 + 24); + float _tmem_load_7[32]; + tmem_ld_x8(&_tmem_load_7[0], o_base_epi + c_1 + 32); + tmem_ld_x8(&_tmem_load_7[8], o_base_epi + c_1 + 32 + 8); + tmem_ld_x8(&_tmem_load_7[16], o_base_epi + c_1 + 32 + 16); + tmem_ld_x8(&_tmem_load_7[24], o_base_epi + c_1 + 32 + 24); + unsigned int epilogue_packed_first[16]; + unsigned int epilogue_packed_next[16]; + const float2 _scale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_6)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_6)[_ls], _scale2_1); + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_6[_lp*2 + 0], _tmem_load_6[_lp*2+1 + 0])); + epilogue_packed_first[_lp] = *(uint32_t*)&_bf2; + } + const float2 _scale2_2 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_7)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_7)[_ls], _scale2_2); + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_7[_lp*2 + 0], _tmem_load_7[_lp*2+1 + 0])); + epilogue_packed_next[_lp] = *(uint32_t*)&_bf2; + } + int gmem_base = o_offset + logical_vs * 64 + c_1; + if (((unsigned long long)O & 31) == 0) { + #pragma unroll + for (int j = 0; j < 32; j += 16) { + { + unsigned _stv8_3_0 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 0]); + unsigned _stv8_3_1 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 1]); + unsigned _stv8_3_2 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 2]); + unsigned _stv8_3_3 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 3]); + unsigned _stv8_3_4 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 4]); + unsigned _stv8_3_5 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 5]); + unsigned _stv8_3_6 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 6]); + unsigned _stv8_3_7 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (gmem_base + j) + (0))), "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3), "r"(_stv8_3_4), "r"(_stv8_3_5), "r"(_stv8_3_6), "r"(_stv8_3_7) : "memory"); + } + } + #pragma unroll + for (int j_1 = 0; j_1 < 32; j_1 += 16) { + { + unsigned _stv8_4_0 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 0]); + unsigned _stv8_4_1 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 1]); + unsigned _stv8_4_2 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 2]); + unsigned _stv8_4_3 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 3]); + unsigned _stv8_4_4 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 4]); + unsigned _stv8_4_5 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 5]); + unsigned _stv8_4_6 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 6]); + unsigned _stv8_4_7 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (gmem_base + 32 + j_1) + (0))), "r"(_stv8_4_0), "r"(_stv8_4_1), "r"(_stv8_4_2), "r"(_stv8_4_3), "r"(_stv8_4_4), "r"(_stv8_4_5), "r"(_stv8_4_6), "r"(_stv8_4_7) : "memory"); + } + } + } else { + #pragma unroll + for (int j_2 = 0; j_2 < 32; j_2 += 8) { + reinterpret_cast(O + (gmem_base + j_2))[0] = reinterpret_cast(epilogue_packed_first + j_2 / 2)[0]; + } + #pragma unroll + for (int j_3 = 0; j_3 < 32; j_3 += 8) { + reinterpret_cast(O + (gmem_base + 32 + j_3))[0] = reinterpret_cast(epilogue_packed_next + j_3 / 2)[0]; + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)(((EXACT007826_CONSTANTS) ? 16 : value_split_count)) + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * ((EXACT007826_CONSTANTS) ? 16 : value_split_count) + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < ((EXACT007826_CONSTANTS) ? 8192 : total_work_items); work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) > 76) { + if (((EXACT007826_CONSTANTS) ? 8192 : total_work_items) <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[((EXACT007826_CONSTANTS) ? 8192 : total_work_items) * ((EXACT007826_CONSTANTS) ? 32 : max_pages_per_seq) + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)(((EXACT007826_CONSTANTS) ? 16 : value_split_count)) + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * ((EXACT007826_CONSTANTS) ? 16 : value_split_count) + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < ((EXACT007826_CONSTANTS) ? 8192 : total_work_items); _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)(((EXACT007826_CONSTANTS) ? 16 : value_split_count)) + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_5f6f81e2ad.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_5f6f81e2ad.cu new file mode 100644 index 00000000000..5f52099bed7 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_5f6f81e2ad.cu @@ -0,0 +1,2433 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < total_work_items; work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (total_work_items > 76) { + if (total_work_items <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * max_pages_per_seq; + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)value_split_count + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < total_work_items; _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)value_split_count + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < total_work_items; work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (total_work_items > 76) { + if (total_work_items <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)value_split_count + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < total_work_items; work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (total_work_items > 76) { + if (total_work_items <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_8[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_8[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_8[j + 0], _tmem_load_8[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_8[j + 2], _tmem_load_8[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_8[j + 4], _tmem_load_8[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_8[j + 6], _tmem_load_8[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)value_split_count + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < total_work_items; work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (total_work_items > 76) { + if (total_work_items <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)value_split_count + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < total_work_items; _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)value_split_count + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_823f60e1b7.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_823f60e1b7.cu new file mode 100644 index 00000000000..9477a951ac7 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_823f60e1b7.cu @@ -0,0 +1,2463 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ uint32_t mbarrier_test_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.test_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_test_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 80;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < total_work_items; work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (total_work_items > 76) { + if (total_work_items <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * max_pages_per_seq; + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)value_split_count + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < total_work_items; _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)value_split_count + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < total_work_items; work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (total_work_items > 76) { + if (total_work_items <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (int)(((unsigned int)seqlen_kv_b_1 + 128 - 1) / 128); + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + uint32_t _mbar_token_10 = mbarrier_test_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + int abs_tile = start_tile_1 + tile_1; + int tile_end_k = (abs_tile + 1) * 128; + bool score_tile_full = tile_end_k <= seqlen_kv_b_1; + float sv[64]; + float tile_max = -MLA_INF; + float new_max = -MLA_INF; + if (!score_tile_full) { + mbarrier_wait_token(s_full_addr + (phase) * 8, s_wait_phase, _mbar_token_10); + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]), "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]), "=f"(tile_max) + : "r"(s_base)); + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + } else { + mbarrier_wait_token(s_full_addr + (phase) * 8, s_wait_phase, _mbar_token_10); + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]), "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]), "=f"(tile_max) + : "r"(s_base)); + new_max = tile_max; + } + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)value_split_count + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 216;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < total_work_items; work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (total_work_items > 76) { + if (total_work_items <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + { + { + #pragma unroll 1 + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll 1 + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + const float2 _scale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_8)[_ls] = mul_f32x2_rn_noftz(reinterpret_cast(_tmem_load_8)[_ls], _scale2_1); + unsigned int pair_packed[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_8[_lp*2 + 0], _tmem_load_8[_lp*2+1 + 0])); + pair_packed[_lp] = *(uint32_t*)&_bf2; + } + int pair_parity = lane & 1; + unsigned int pair_mask = (unsigned int)(-pair_parity); + int pair_neighbor = lane ^ 1; + int pair_even_base = gmem_base - pair_parity * 512 + pair_parity * 8; + #pragma unroll + for (int pair_columns = 0; pair_columns < 2; pair_columns++) { + unsigned int pair_even_words[4]; + unsigned int pair_odd_words[4]; + #pragma unroll + for (int pair_word = 0; pair_word < 4; pair_word++) { + unsigned int _shfl_0 = __shfl_sync(0xFFFFFFFF, pair_packed[pair_columns * 8 + 4 + pair_word] & ~pair_mask | pair_packed[pair_columns * 8 + pair_word] & pair_mask, pair_neighbor); + pair_even_words[pair_word] = pair_packed[pair_columns * 8 + pair_word] & ~pair_mask | _shfl_0 & pair_mask; + pair_odd_words[pair_word] = _shfl_0 & ~pair_mask | pair_packed[pair_columns * 8 + 4 + pair_word] & pair_mask; + } + reinterpret_cast(O + (pair_even_base + pair_columns * 16))[0] = reinterpret_cast(pair_even_words)[0]; + reinterpret_cast(O + (pair_even_base + 512 + pair_columns * 16))[0] = reinterpret_cast(pair_odd_words)[0]; + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)value_split_count + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < total_work_items; work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (total_work_items > 76) { + if (total_work_items <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)value_split_count + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < total_work_items; _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)value_split_count + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_a6a70d7004.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_a6a70d7004.cu new file mode 100644 index 00000000000..d800bf6e5e1 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_a6a70d7004.cu @@ -0,0 +1,2434 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 +#define Q8_B768 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < ((Q8_B768) ? 6144 : total_work_items); work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 76) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * ((Q8_B768) ? 32 : max_pages_per_seq); + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + if (n == 0) { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + if (n_1 == 0) { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < ((Q8_B768) ? 6144 : total_work_items); _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < ((Q8_B768) ? 6144 : total_work_items); work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 76) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]), "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]), "=f"(tile_max) + : "r"(s_base)); + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + if (!score_tile_full) { + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + } + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < ((Q8_B768) ? 6144 : total_work_items); work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 76) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll 1 + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_8[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_8[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_8[j + 0], _tmem_load_8[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_8[j + 2], _tmem_load_8[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_8[j + 4], _tmem_load_8[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_8[j + 6], _tmem_load_8[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < ((Q8_B768) ? 6144 : total_work_items); work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 76) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < ((Q8_B768) ? 6144 : total_work_items); _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_d003c22422.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_d003c22422.cu new file mode 100644 index 00000000000..5f7b2eae798 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_d003c22422.cu @@ -0,0 +1,2503 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 +#define Q4_B1024 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ void tmem_ld_x8(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x8.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * ((Q4_B1024) ? 4 : value_split_count) + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < ((Q4_B1024) ? 4096 : total_work_items); work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (((Q4_B1024) ? 4096 : total_work_items) > 76) { + if (((Q4_B1024) ? 4096 : total_work_items) <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[((Q4_B1024) ? 4096 : total_work_items) * ((Q4_B1024) ? 32 : max_pages_per_seq) + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * ((Q4_B1024) ? 32 : max_pages_per_seq); + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)(((Q4_B1024) ? 4 : value_split_count)) + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * ((Q4_B1024) ? 4 : value_split_count) + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < ((Q4_B1024) ? 4096 : total_work_items); _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)(((Q4_B1024) ? 4 : value_split_count)) + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 176;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * ((Q4_B1024) ? 4 : value_split_count) + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < ((Q4_B1024) ? 4096 : total_work_items); work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (((Q4_B1024) ? 4096 : total_work_items) > 76) { + if (((Q4_B1024) ? 4096 : total_work_items) <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[((Q4_B1024) ? 4096 : total_work_items) * ((Q4_B1024) ? 32 : max_pages_per_seq) + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]), "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]), "=f"(tile_max) + : "r"(s_base)); + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + if (!score_tile_full) { + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + } + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + float _tmem_load_0[1]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x1.b32" + " {%0}, [%1];" + : "=f"(_tmem_load_0[0]) + : "r"(exact_meta_addr + 1)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + new_max = _tmem_load_0[0]; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)(((Q4_B1024) ? 4 : value_split_count)) + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * ((Q4_B1024) ? 4 : value_split_count) + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < ((Q4_B1024) ? 4096 : total_work_items); work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (((Q4_B1024) ? 4096 : total_work_items) > 76) { + if (((Q4_B1024) ? 4096 : total_work_items) <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[((Q4_B1024) ? 4096 : total_work_items) * ((Q4_B1024) ? 32 : max_pages_per_seq) + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_1[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_1[0] - _tmem_load_1[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_1[1] <= _tmem_load_1[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_3[16]; + tmem_ld_x16(&_tmem_load_3[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_3)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_3)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_3); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_5[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_5[0]; + final_max_val = _tmem_load_5[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll 1 + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll 1 + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 64) { + float _tmem_load_7[32]; + tmem_ld_x8(&_tmem_load_7[0], o_base_epi + c_1); + tmem_ld_x8(&_tmem_load_7[8], o_base_epi + c_1 + 8); + tmem_ld_x8(&_tmem_load_7[16], o_base_epi + c_1 + 16); + tmem_ld_x8(&_tmem_load_7[24], o_base_epi + c_1 + 24); + float _tmem_load_8[32]; + tmem_ld_x8(&_tmem_load_8[0], o_base_epi + c_1 + 32); + tmem_ld_x8(&_tmem_load_8[8], o_base_epi + c_1 + 32 + 8); + tmem_ld_x8(&_tmem_load_8[16], o_base_epi + c_1 + 32 + 16); + tmem_ld_x8(&_tmem_load_8[24], o_base_epi + c_1 + 32 + 24); + unsigned int epilogue_packed_first[16]; + unsigned int epilogue_packed_next[16]; + const float2 _scale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_7)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_7)[_ls], _scale2_1); + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_7[_lp*2 + 0], _tmem_load_7[_lp*2+1 + 0])); + epilogue_packed_first[_lp] = *(uint32_t*)&_bf2; + } + const float2 _scale2_2 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_8)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_8)[_ls], _scale2_2); + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_8[_lp*2 + 0], _tmem_load_8[_lp*2+1 + 0])); + epilogue_packed_next[_lp] = *(uint32_t*)&_bf2; + } + int gmem_base = o_offset + logical_vs * 64 + c_1; + if (((unsigned long long)O & 31) == 0) { + #pragma unroll + for (int j = 0; j < 32; j += 16) { + { + unsigned _stv8_3_0 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 0]); + unsigned _stv8_3_1 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 1]); + unsigned _stv8_3_2 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 2]); + unsigned _stv8_3_3 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 3]); + unsigned _stv8_3_4 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 4]); + unsigned _stv8_3_5 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 5]); + unsigned _stv8_3_6 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 6]); + unsigned _stv8_3_7 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (gmem_base + j) + (0))), "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3), "r"(_stv8_3_4), "r"(_stv8_3_5), "r"(_stv8_3_6), "r"(_stv8_3_7) : "memory"); + } + } + #pragma unroll + for (int j_1 = 0; j_1 < 32; j_1 += 16) { + { + unsigned _stv8_4_0 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 0]); + unsigned _stv8_4_1 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 1]); + unsigned _stv8_4_2 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 2]); + unsigned _stv8_4_3 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 3]); + unsigned _stv8_4_4 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 4]); + unsigned _stv8_4_5 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 5]); + unsigned _stv8_4_6 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 6]); + unsigned _stv8_4_7 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (gmem_base + 32 + j_1) + (0))), "r"(_stv8_4_0), "r"(_stv8_4_1), "r"(_stv8_4_2), "r"(_stv8_4_3), "r"(_stv8_4_4), "r"(_stv8_4_5), "r"(_stv8_4_6), "r"(_stv8_4_7) : "memory"); + } + } + } else { + #pragma unroll + for (int j_2 = 0; j_2 < 32; j_2 += 8) { + reinterpret_cast(O + (gmem_base + j_2))[0] = reinterpret_cast(epilogue_packed_first + j_2 / 2)[0]; + } + #pragma unroll + for (int j_3 = 0; j_3 < 32; j_3 += 8) { + reinterpret_cast(O + (gmem_base + 32 + j_3))[0] = reinterpret_cast(epilogue_packed_next + j_3 / 2)[0]; + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)(((Q4_B1024) ? 4 : value_split_count)) + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * ((Q4_B1024) ? 4 : value_split_count) + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < ((Q4_B1024) ? 4096 : total_work_items); work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (((Q4_B1024) ? 4096 : total_work_items) > 76) { + if (((Q4_B1024) ? 4096 : total_work_items) <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[((Q4_B1024) ? 4096 : total_work_items) * ((Q4_B1024) ? 32 : max_pages_per_seq) + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)(((Q4_B1024) ? 4 : value_split_count)) + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * ((Q4_B1024) ? 4 : value_split_count) + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < ((Q4_B1024) ? 4096 : total_work_items); _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)(((Q4_B1024) ? 4 : value_split_count)) + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_d3892bb3e6.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_d3892bb3e6.cu new file mode 100644 index 00000000000..187ab11c1ca --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_d3892bb3e6.cu @@ -0,0 +1,2479 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ uint32_t mbarrier_test_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.test_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_test_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < total_work_items; work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (total_work_items > 76) { + if (total_work_items <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * max_pages_per_seq; + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)value_split_count + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < total_work_items; _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)value_split_count + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < total_work_items; work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (total_work_items > 76) { + if (total_work_items <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + unsigned int next_score_ready_token = 0; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait_token(s_full_addr + (phase) * 8, s_wait_phase, next_score_ready_token); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + float tile_max_low = -MLA_INF; + float tile_max_high = -MLA_INF; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x32.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, %32, [%33];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]), "=f"(tile_max_low) + : "r"(s_base)); + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x32.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, %32, [%33];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]), "=f"(tile_max_high) + : "r"(s_base + 32)); + float _max_0 = max_noftz(tile_max_low, tile_max_high); + tile_max = _max_0; + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + if (!score_tile_full) { + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + } + float _max_1 = max_noftz(new_max, row_max_val); + new_max = _max_1; + if (num_kv_tiles_1 > tile_1 + 1) { + uint32_t _mbar_token_10 = mbarrier_test_wait(s_full_addr + (pipeline_tile + 1 & 1) * 8, pipeline_tile + 1 >> 1 & 1); + next_score_ready_token = _mbar_token_10; + } + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_2 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_2; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)value_split_count + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < total_work_items; work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (total_work_items > 76) { + if (total_work_items <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_8[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_8[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_8[j + 0], _tmem_load_8[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_8[j + 2], _tmem_load_8[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_8[j + 4], _tmem_load_8[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_8[j + 6], _tmem_load_8[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)value_split_count + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < total_work_items; work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (total_work_items > 76) { + if (total_work_items <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[total_work_items * max_pages_per_seq + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)value_split_count + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * value_split_count + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < total_work_items; _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)value_split_count + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_d56d2460cb.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_d56d2460cb.cu new file mode 100644 index 00000000000..71599316186 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_d56d2460cb.cu @@ -0,0 +1,2454 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 +#define Q1_B512 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * ((Q1_B512) ? 1 : value_split_count) + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < ((Q1_B512) ? 512 : total_work_items); work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (((Q1_B512) ? 512 : total_work_items) > 76) { + if (((Q1_B512) ? 512 : total_work_items) <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[((Q1_B512) ? 512 : total_work_items) * max_pages_per_seq + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * max_pages_per_seq; + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)(((Q1_B512) ? 1 : value_split_count)) + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * ((Q1_B512) ? 1 : value_split_count) + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < ((Q1_B512) ? 512 : total_work_items); _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)(((Q1_B512) ? 1 : value_split_count)) + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * ((Q1_B512) ? 1 : value_split_count) + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < ((Q1_B512) ? 512 : total_work_items); work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (((Q1_B512) ? 512 : total_work_items) > 76) { + if (((Q1_B512) ? 512 : total_work_items) <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[((Q1_B512) ? 512 : total_work_items) * max_pages_per_seq + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + _reg_reduce_max2_2.x = max_noftz(_reg_reduce_max2_2.x, max_noftz(sv[0], sv[1])); + _reg_reduce_max2_2.y = max_noftz(_reg_reduce_max2_2.y, max_noftz(sv[2], sv[3])); + _reg_reduce_max2_2.x = max_noftz(_reg_reduce_max2_2.x, max_noftz(sv[4], sv[5])); + _reg_reduce_max2_2.y = max_noftz(_reg_reduce_max2_2.y, max_noftz(sv[6], sv[7])); + _reg_reduce_max2_2.x = max_noftz(_reg_reduce_max2_2.x, max_noftz(sv[8], sv[9])); + _reg_reduce_max2_2.y = max_noftz(_reg_reduce_max2_2.y, max_noftz(sv[10], sv[11])); + _reg_reduce_max2_2.x = max_noftz(_reg_reduce_max2_2.x, max_noftz(sv[12], sv[13])); + _reg_reduce_max2_2.y = max_noftz(_reg_reduce_max2_2.y, max_noftz(sv[14], sv[15])); + _reg_reduce_max2_2.x = max_noftz(_reg_reduce_max2_2.x, max_noftz(sv[16], sv[17])); + _reg_reduce_max2_2.y = max_noftz(_reg_reduce_max2_2.y, max_noftz(sv[18], sv[19])); + _reg_reduce_max2_2.x = max_noftz(_reg_reduce_max2_2.x, max_noftz(sv[20], sv[21])); + _reg_reduce_max2_2.y = max_noftz(_reg_reduce_max2_2.y, max_noftz(sv[22], sv[23])); + _reg_reduce_max2_2.x = max_noftz(_reg_reduce_max2_2.x, max_noftz(sv[24], sv[25])); + _reg_reduce_max2_2.y = max_noftz(_reg_reduce_max2_2.y, max_noftz(sv[26], sv[27])); + _reg_reduce_max2_2.x = max_noftz(_reg_reduce_max2_2.x, max_noftz(sv[28], sv[29])); + _reg_reduce_max2_2.y = max_noftz(_reg_reduce_max2_2.y, max_noftz(sv[30], sv[31])); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + float max_low = sv_max; + float2 _reg_reduce_max2_3 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[(32) + 0], _reg_reduce_max2_3); + float sv_max_0 = row_max_reduce(_reg_reduce_max2_3); + float max_high = sv_max_0; + float _max_0 = max_noftz(max_low, max_high); + new_max = _max_0; + float _max_1 = max_noftz(new_max, row_max_val); + new_max = _max_1; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_2 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_2; + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + new_max = new_max; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_4 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_4); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_4); + float sv_sum = _reg_reduce_sum2_4.x + _reg_reduce_sum2_4.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)(((Q1_B512) ? 1 : value_split_count)) + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * ((Q1_B512) ? 1 : value_split_count) + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < ((Q1_B512) ? 512 : total_work_items); work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (((Q1_B512) ? 512 : total_work_items) > 76) { + if (((Q1_B512) ? 512 : total_work_items) <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[((Q1_B512) ? 512 : total_work_items) * max_pages_per_seq + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_0[0] - _tmem_load_0[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, _tmem_load_0[1] <= _tmem_load_0[0]); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_4[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_4[0]; + final_max_val = _tmem_load_4[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_8[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_8[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_8[j + 0], _tmem_load_8[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_8[j + 2], _tmem_load_8[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_8[j + 4], _tmem_load_8[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_8[j + 6], _tmem_load_8[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)(((Q1_B512) ? 1 : value_split_count)) + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * ((Q1_B512) ? 1 : value_split_count) + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < ((Q1_B512) ? 512 : total_work_items); work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (((Q1_B512) ? 512 : total_work_items) > 76) { + if (((Q1_B512) ? 512 : total_work_items) <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[((Q1_B512) ? 512 : total_work_items) * max_pages_per_seq + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)(((Q1_B512) ? 1 : value_split_count)) + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * ((Q1_B512) ? 1 : value_split_count) + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < ((Q1_B512) ? 512 : total_work_items); _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)(((Q1_B512) ? 1 : value_split_count)) + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_fb6e96f4d8.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_fb6e96f4d8.cu new file mode 100644 index 00000000000..5829e6a07a0 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_exact_live_bf16_clc_fb6e96f4d8.cu @@ -0,0 +1,2525 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 73728 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 73728 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 73728 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 229376 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 229888 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_WORK_RESPONSE_OFF 230400 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 230912 +#define THREADS 384 +#define Q8_B768 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_mla_decode_exact_live_bf16_clc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem + 230432; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + #define work_full_addr (mbar_base + 384) + #define work_empty_addr (mbar_base + 400) + #define throttle_full_addr (mbar_base + 416) + #define throttle_empty_addr (mbar_base + 432) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_kv_addr = smem + 73728; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vhalf_addr = smem + 73728; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 73728); + const int smem_vquarter_addr = smem + 73728; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 229376); + const int smem_softmax_exchange_addr = smem + 229376; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 229888); + const int smem_epilogue_exchange_addr = smem + 229888; + unsigned int* work_response = reinterpret_cast(smem_raw + 230400); + const int work_response_addr = smem + 230400; + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers) + // Mbarriers at smem_raw[230432..230880) + + if (warp == 0) { + // q_full: 1 barriers, init_count=1 + // q_empty: 1 barriers, init_count=1 + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + // kv_empty: 15 barriers, init_count=1 + // s_full: 2 barriers, init_count=1 + // s_empty: 2 barriers, init_count=256 + // p_full: 2 barriers, init_count=256 + // p_empty: 2 barriers, init_count=1 + // o_empty: 1 barriers, init_count=256 + // stats: 2 barriers, init_count=128 + // stats_empty: 2 barriers, init_count=128 + // o_full: 1 barriers, init_count=1 + // tmem_scrubbed: 1 barriers, init_count=256 + // tmem_dealloc_peer: 1 barriers, init_count=32 + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + // work_empty: 2 barriers, init_count=704 + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=32 + // throttle_empty: 2 barriers, init_count=32 + // Warp-cooperative initialization in physical record order. + mbarrier_init(smem + 230432 + lane * 8, 1); + uint32_t _mbarrier_init_count_0_32 = 32; + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(20), "r"((uint32_t)(704))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(18), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(16), "r"((uint32_t)(32))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(14), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(128))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(9), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(8), "r"((uint32_t)(1))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(6), "r"((uint32_t)(256))); + asm volatile("{ .reg .pred p; setp.lt.u32 p, %1, %2; selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(2), "r"((uint32_t)(1))); + if (lane < 24) { + mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230880); + if (warp == 0) { + int _tmem_hold = smem + 230880; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + // Partial post-allocation TMEM rendezvous (384 threads on named barrier 0) + if (warp <= 11) { + asm volatile("barrier.sync.aligned %0, %1;" :: "r"(0), "r"(384) : "memory"); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + + const int taddr = (warp <= 11) ? tmem_addr_storage[0] : 0; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int load_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int load_work_stage = 0; + unsigned int load_throttle_stage = 0; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_idx = 0; work_idx < ((Q8_B768) ? 6144 : total_work_items); work_idx++) { + unsigned int load_ordered_work_idx = load_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 76) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + load_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)load_work_idx]; + } + } + if (cta_rank == 0) { + mbarrier_wait(throttle_empty_addr + (load_throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (load_throttle_stage) * 8); + load_throttle_stage += 1; + if (load_throttle_stage == 2) { load_throttle_stage = 0; _phase_throttle_empty ^= 1; } + } + int batch_idx = load_ordered_work_idx; + int value_split = 0; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + int pt_base = batch_idx * ((Q8_B768) ? 32 : max_pages_per_seq); + int q_row_global = batch_idx * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n = 0; n < 8; n++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, (&tmap_k), 0, 0, n, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k), 0, 0, n, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + int abs_tile_k = start_tile + tile; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k), 0, 0, n_1, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k), 0, 0, n_1, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile + tile - 1; + #pragma unroll + for (int ps = 0; ps < 4; ps++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps]; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((1) ? 2 : 1); acc_stage++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, (&tmap_v), 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, (&tmap_v), 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile + num_kv_tiles - 1; + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_1]; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((1) ? 2 : 1); acc_stage_1++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_1 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, (&tmap_v), 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, (&tmap_v), 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_0 = 0; + uint32_t _clc_ctaid_x_0; + uint32_t _clc_ctaid_y_0; + uint32_t _clc_ctaid_z_0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_0), "=r"(_clc_ctaid_y_0), "=r"(_clc_ctaid_z_0), "=r"(_clc_valid_0) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + load_work_stage * 8), "r"(0) : "memory"); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full ^= 1; } + unsigned int load_work_valid = _clc_valid_0; + if (load_work_valid == 0) { + break; + } + load_work_idx = _clc_ctaid_z_0 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_0 / 2; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + // ---- Role: empty1 ---- + } else if (warp == 11) { + { // empty1_main + const int wg2_dummy_1 = 0; + unsigned int empty1_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int empty1_work_stage = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_iter = 0; _work_iter < ((Q8_B768) ? 6144 : total_work_items); _work_iter++) { + mbarrier_wait(work_full_addr + (empty1_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_5 = 0; + uint32_t _clc_ctaid_x_5; + uint32_t _clc_ctaid_y_5; + uint32_t _clc_ctaid_z_5; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_5), "=r"(_clc_ctaid_y_5), "=r"(_clc_ctaid_z_5), "=r"(_clc_valid_5) + : "r"(work_response_addr + empty1_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + empty1_work_stage * 8), "r"(0) : "memory"); + empty1_work_stage += 1; + if (empty1_work_stage == 2) { empty1_work_stage = 0; _phase_work_full_1 ^= 1; } + unsigned int empty1_work_valid = _clc_valid_5; + if (empty1_work_valid == 0) { + break; + } + empty1_work_idx = _clc_ctaid_z_5 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_5 / 2; + } + } + // ---- Role: softmax_wg ---- + } else if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + unsigned int softmax_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int softmax_work_stage = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = 0; work_idx_1 < ((Q8_B768) ? 6144 : total_work_items); work_idx_1++) { + unsigned int softmax_ordered_work_idx = softmax_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 76) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + softmax_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)softmax_work_idx]; + } + } + int batch_idx_1 = softmax_ordered_work_idx; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + float true_row_max_log2 = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]), "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]), "=f"(tile_max) + : "r"(s_base)); + int abs_tile = start_tile_1 + tile_1; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b_1 % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b_1; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b_1 - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + if (!score_tile_full) { + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + } + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + { + float candidate_max_log2 = new_max * softmax_scale_log2; + float normalization_max_log2 = row_max_val * softmax_scale_log2; + float _max_2 = max_noftz(true_row_max_log2, candidate_max_log2); + true_row_max_log2 = _max_2; + int _vote_0 = __any_sync(0xFFFFFFFF, candidate_max_log2 - normalization_max_log2 > 6.0f); + int rescale_vote = _vote_0; + if (rescale_vote == 0) { + new_max = row_max_val; + } + } + float exact_meta[4]; + int exact_meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + mbarrier_wait(stats_empty_addr + (phase) * 8, (unsigned int)(pipeline_tile >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + exact_meta[0] = row_max_val; + exact_meta[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(exact_meta_addr), "f"(exact_meta[0]), "f"(exact_meta[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + float _tmem_load_0[1]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x1.b32" + " {%0}, [%1];" + : "=f"(_tmem_load_0[0]) + : "r"(exact_meta_addr + 1)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + new_max = _tmem_load_0[0]; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + float meta[4]; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps_2 = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps_2; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + seed_zero[0] = row_sum_val; + seed_zero[1] = row_max_val; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + 384 + (unsigned int)((softmax_tile_cursor + num_kv_tiles_1 & 1) * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16)), "f"(seed_zero[0]), "f"(seed_zero[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 & 1) * 8); + mbarrier_wait(stats_empty_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8, (unsigned int)(softmax_tile_cursor + num_kv_tiles_1 + 1 >> 1 & 1) ^ softmax_work_stage & 1 ^ 1); + mbarrier_arrive(stats_addr + (softmax_tile_cursor + num_kv_tiles_1 + 1 & 1) * 8); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + uint32_t _clc_ctaid_x_2; + uint32_t _clc_ctaid_y_2; + uint32_t _clc_ctaid_z_2; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_2), "=r"(_clc_ctaid_y_2), "=r"(_clc_ctaid_z_2), "=r"(_clc_valid_2) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + softmax_work_stage * 8), "r"(0) : "memory"); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full_2 ^= 1; } + unsigned int softmax_work_valid = _clc_valid_2; + if (softmax_work_valid == 0) { + break; + } + softmax_work_idx = _clc_ctaid_z_2 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_2 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + // ---- Role: correction_wg ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int correction_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int correction_work_stage = 0; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = 0; work_idx_2 < ((Q8_B768) ? 6144 : total_work_items); work_idx_2++) { + unsigned int correction_ordered_work_idx = correction_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 76) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + correction_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)correction_work_idx]; + } + } + int batch_idx_2 = correction_ordered_work_idx; + int value_split_1 = 0; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + asm volatile("griddepcontrol.wait;" ::: "memory"); + mbarrier_wait(stats_addr + (correction_tile_cursor & 1) * 8, (unsigned int)(correction_tile_cursor >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor & 1) * 8); + #pragma unroll 1 + for (int tile_2 = 1; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = (unsigned int)(pipeline_tile_1 >> 1 & 1) ^ correction_work_stage & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_1[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]) + : "r"(meta_addr_1)); + final_sum_val = final_sum_val; + final_max_val = final_max_val; + float _exp2_3 = approx_exp2(softmax_scale_log2 * (_tmem_load_1[0] - _tmem_load_1[1])); + float acc_scale_1 = _exp2_3; + float no_correction_1 = 0.0f; + { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + int _vote_1 = __all_sync(0xFFFFFFFF, _tmem_load_1[1] <= _tmem_load_1[0]); + int skip_correction = _vote_1; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage_2 * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + float _tmem_load_3[16]; + tmem_ld_x16(&_tmem_load_3[0], o_base); + float _tmem_load_4[16]; + tmem_ld_x16(&_tmem_load_4[0], o_base + 16); + float _tmem_load_5[16]; + tmem_ld_x16(&_tmem_load_5[0], o_base + 32); + float _tmem_load_6[16]; + tmem_ld_x16(&_tmem_load_6[0], o_base + 48); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_3)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_3)[_ls], _scale2_0); + tmem_st_x16_f32(o_base, _tmem_load_3); + const float2 _scale2_1 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_4)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_4)[_ls], _scale2_1); + tmem_st_x16_f32(o_base + 16, _tmem_load_4); + const float2 _scale2_2 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_5)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_5)[_ls], _scale2_2); + tmem_st_x16_f32(o_base + 32, _tmem_load_5); + const float2 _scale2_3 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_6)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_6)[_ls], _scale2_3); + tmem_st_x16_f32(o_base + 48, _tmem_load_6); + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 >> 1 & 1) ^ correction_work_stage & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_8[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]) + : "r"(taddr + 384 + (unsigned int)((correction_tile_cursor + num_kv_tiles_2 & 1) * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row)); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 & 1) * 8); + mbarrier_wait(stats_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8, (unsigned int)(correction_tile_cursor + num_kv_tiles_2 + 1 >> 1 & 1) ^ correction_work_stage & 1); + mbarrier_arrive(stats_empty_addr + (correction_tile_cursor + num_kv_tiles_2 + 1 & 1) * 8); + final_sum_val = _tmem_load_8[0]; + final_max_val = _tmem_load_8[1]; + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_2 * 128 + head_idx) * 512; + float output_scale = 1.0f; + { + { + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_3 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_3 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c = 0; c < 64; c += 64) { + float _tmem_load_10[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_10[0]), "=f"(_tmem_load_10[1]), "=f"(_tmem_load_10[2]), "=f"(_tmem_load_10[3]), "=f"(_tmem_load_10[4]), "=f"(_tmem_load_10[5]), "=f"(_tmem_load_10[6]), "=f"(_tmem_load_10[7]), "=f"(_tmem_load_10[8]), "=f"(_tmem_load_10[9]), "=f"(_tmem_load_10[10]), "=f"(_tmem_load_10[11]), "=f"(_tmem_load_10[12]), "=f"(_tmem_load_10[13]), "=f"(_tmem_load_10[14]), "=f"(_tmem_load_10[15]), "=f"(_tmem_load_10[16]), "=f"(_tmem_load_10[17]), "=f"(_tmem_load_10[18]), "=f"(_tmem_load_10[19]), "=f"(_tmem_load_10[20]), "=f"(_tmem_load_10[21]), "=f"(_tmem_load_10[22]), "=f"(_tmem_load_10[23]), "=f"(_tmem_load_10[24]), "=f"(_tmem_load_10[25]), "=f"(_tmem_load_10[26]), "=f"(_tmem_load_10[27]), "=f"(_tmem_load_10[28]), "=f"(_tmem_load_10[29]), "=f"(_tmem_load_10[30]), "=f"(_tmem_load_10[31]) + : "r"(o_base_epi + c)); + float _tmem_load_11[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_11[0]), "=f"(_tmem_load_11[1]), "=f"(_tmem_load_11[2]), "=f"(_tmem_load_11[3]), "=f"(_tmem_load_11[4]), "=f"(_tmem_load_11[5]), "=f"(_tmem_load_11[6]), "=f"(_tmem_load_11[7]), "=f"(_tmem_load_11[8]), "=f"(_tmem_load_11[9]), "=f"(_tmem_load_11[10]), "=f"(_tmem_load_11[11]), "=f"(_tmem_load_11[12]), "=f"(_tmem_load_11[13]), "=f"(_tmem_load_11[14]), "=f"(_tmem_load_11[15]), "=f"(_tmem_load_11[16]), "=f"(_tmem_load_11[17]), "=f"(_tmem_load_11[18]), "=f"(_tmem_load_11[19]), "=f"(_tmem_load_11[20]), "=f"(_tmem_load_11[21]), "=f"(_tmem_load_11[22]), "=f"(_tmem_load_11[23]), "=f"(_tmem_load_11[24]), "=f"(_tmem_load_11[25]), "=f"(_tmem_load_11[26]), "=f"(_tmem_load_11[27]), "=f"(_tmem_load_11[28]), "=f"(_tmem_load_11[29]), "=f"(_tmem_load_11[30]), "=f"(_tmem_load_11[31]) + : "r"(o_base_epi + c + 32)); + unsigned int epilogue_packed_first[16]; + unsigned int epilogue_packed_next[16]; + const float2 _scale2_4 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_10)[_ls] = mul_f32x2_rn_noftz(reinterpret_cast(_tmem_load_10)[_ls], _scale2_4); + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_10[_lp*2 + 0], _tmem_load_10[_lp*2+1 + 0])); + epilogue_packed_first[_lp] = *(uint32_t*)&_bf2; + } + const float2 _scale2_5 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_11)[_ls] = mul_f32x2_rn_noftz(reinterpret_cast(_tmem_load_11)[_ls], _scale2_5); + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_11[_lp*2 + 0], _tmem_load_11[_lp*2+1 + 0])); + epilogue_packed_next[_lp] = *(uint32_t*)&_bf2; + } + int gmem_base = o_offset + logical_vs * 64 + c; + if (((unsigned long long)O & 31) == 0) { + #pragma unroll + for (int j = 0; j < 32; j += 16) { + { + unsigned _stv8_6_0 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 0]); + unsigned _stv8_6_1 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 1]); + unsigned _stv8_6_2 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 2]); + unsigned _stv8_6_3 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 3]); + unsigned _stv8_6_4 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 4]); + unsigned _stv8_6_5 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 5]); + unsigned _stv8_6_6 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 6]); + unsigned _stv8_6_7 = __float_as_uint(reinterpret_cast(epilogue_packed_first)[j / 2 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (gmem_base + j) + (0))), "r"(_stv8_6_0), "r"(_stv8_6_1), "r"(_stv8_6_2), "r"(_stv8_6_3), "r"(_stv8_6_4), "r"(_stv8_6_5), "r"(_stv8_6_6), "r"(_stv8_6_7) : "memory"); + } + } + #pragma unroll + for (int j_1 = 0; j_1 < 32; j_1 += 16) { + { + unsigned _stv8_7_0 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 0]); + unsigned _stv8_7_1 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 1]); + unsigned _stv8_7_2 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 2]); + unsigned _stv8_7_3 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 3]); + unsigned _stv8_7_4 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 4]); + unsigned _stv8_7_5 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 5]); + unsigned _stv8_7_6 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 6]); + unsigned _stv8_7_7 = __float_as_uint(reinterpret_cast(epilogue_packed_next)[j_1 / 2 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (gmem_base + 32 + j_1) + (0))), "r"(_stv8_7_0), "r"(_stv8_7_1), "r"(_stv8_7_2), "r"(_stv8_7_3), "r"(_stv8_7_4), "r"(_stv8_7_5), "r"(_stv8_7_6), "r"(_stv8_7_7) : "memory"); + } + } + } else { + #pragma unroll + for (int j_2 = 0; j_2 < 32; j_2 += 8) { + reinterpret_cast(O + (gmem_base + j_2))[0] = reinterpret_cast(epilogue_packed_first + j_2 / 2)[0]; + } + #pragma unroll + for (int j_3 = 0; j_3 < 32; j_3 += 8) { + reinterpret_cast(O + (gmem_base + 32 + j_3))[0] = reinterpret_cast(epilogue_packed_next + j_3 / 2)[0]; + } + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_3 = 0; + uint32_t _clc_ctaid_x_3; + uint32_t _clc_ctaid_y_3; + uint32_t _clc_ctaid_z_3; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_3), "=r"(_clc_ctaid_y_3), "=r"(_clc_ctaid_z_3), "=r"(_clc_valid_3) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + correction_work_stage * 8), "r"(0) : "memory"); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_3 ^= 1; } + unsigned int correction_work_valid = _clc_valid_3; + if (correction_work_valid == 0) { + break; + } + correction_work_idx = _clc_ctaid_z_3 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_3 / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + // ---- Role: mma_warp ---- + } else if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); + const int wg2_dummy_2 = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int mma_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int mma_work_stage = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = 0; work_idx_3 < ((Q8_B768) ? 6144 : total_work_items); work_idx_3++) { + unsigned int mma_ordered_work_idx = mma_work_idx; + if (((Q8_B768) ? 6144 : total_work_items) > 76) { + if (((Q8_B768) ? 6144 : total_work_items) <= 1024) { + mma_ordered_work_idx = (unsigned int)page_table[((Q8_B768) ? 6144 : total_work_items) * ((Q8_B768) ? 32 : max_pages_per_seq) + (int)mma_work_idx]; + } + } + if (cta_rank != 0) { + break; + } + int batch_idx_3 = mma_ordered_work_idx; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n_2) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n_2 == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int p_stage_1 = prev_phase * 4 + ps_3; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_4 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_3 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int p_stage_last = last_phase * 4 + ps_4; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_5 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_4 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_1 = 0; + uint32_t _clc_ctaid_x_1; + uint32_t _clc_ctaid_y_1; + uint32_t _clc_ctaid_z_1; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_1), "=r"(_clc_ctaid_y_1), "=r"(_clc_ctaid_z_1), "=r"(_clc_valid_1) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + mma_work_stage * 8), "r"(0) : "memory"); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_4 ^= 1; } + unsigned int mma_work_valid = _clc_valid_1; + if (mma_work_valid == 0) { + break; + } + mma_work_idx = _clc_ctaid_z_1 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_1 / 2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + // ---- Role: empty0 ---- + } else if (warp == 10) { + { // empty0_main + const int wg2_dummy_3 = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_5 = 0; + if (cta_rank == 0) { + unsigned int scheduler_work_idx = blockIdx.z * ((Q8_B768) ? 8 : value_split_count) + blockIdx.x / 2; + unsigned int scheduler_prod_stage = 0; + unsigned int scheduler_cons_stage = 0; + unsigned int scheduler_throttle_stage = 0; + #pragma unroll 1 + for (unsigned int _work_iter_1 = 0; _work_iter_1 < ((Q8_B768) ? 6144 : total_work_items); _work_iter_1++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + if (lane < 2) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + scheduler_prod_stage * 8), "r"(lane), "r"((uint32_t)(16)) : "memory"); + } + if (elect_sync()) { + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.multicast::cluster::all.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_prod_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_prod_stage * 8) + : "memory"); + } + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + mbarrier_wait(work_full_addr + (scheduler_cons_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_4 = 0; + uint32_t _clc_ctaid_x_4; + uint32_t _clc_ctaid_y_4; + uint32_t _clc_ctaid_z_4; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_x_4), "=r"(_clc_ctaid_y_4), "=r"(_clc_ctaid_z_4), "=r"(_clc_valid_4) + : "r"(work_response_addr + scheduler_cons_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + scheduler_cons_stage * 8), "r"(0) : "memory"); + scheduler_cons_stage += 1; + if (scheduler_cons_stage == 2) { scheduler_cons_stage = 0; _phase_work_full_5 ^= 1; } + unsigned int scheduler_work_valid = _clc_valid_4; + if (scheduler_work_valid == 0) { + break; + } + scheduler_work_idx = _clc_ctaid_z_4 * (unsigned int)(((Q8_B768) ? 8 : value_split_count)) + _clc_ctaid_x_4 / 2; + } + #pragma unroll + for (int _tail = 0; _tail < 2; _tail++) { + mbarrier_wait(work_empty_addr + (scheduler_prod_stage) * 8, _phase_work_empty); + scheduler_prod_stage += 1; + if (scheduler_prod_stage == 2) { scheduler_prod_stage = 0; _phase_work_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_v32_paged_bf16_c80238a3c3.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_v32_paged_bf16_c80238a3c3.cu new file mode 100644 index 00000000000..a0b557ba072 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_v32_paged_bf16_c80238a3c3.cu @@ -0,0 +1,2072 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 74752 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_V_OFF 74752 +#define SMEM_SMEM_V_STAGE_BYTES 8192 +#define SMEM_SMEM_V_STRIDE 8192 +#define SMEM_SMEM_P_OFF 197632 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 230400 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 230912 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_TOTAL 231424 +#define THREADS 384 +#define USE_PDL 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384) __cluster_dims__(2,1,1) void +kernel_mla_decode_v32_paged_bf16(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_k_nope, const __grid_constant__ CUtensorMap tmap_k_rope, const __grid_constant__ CUtensorMap tmap_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int num_split, int total_work_items, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_kv_addr = smem + 74752; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_v_addr = smem + 74752; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 197632); + const int smem_p_addr = smem + 197632; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 230400); + const int smem_softmax_exchange_addr = smem + 230400; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 230912); + const int smem_epilogue_exchange_addr = smem + 230912; + if (warp == 8) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); } + if (warp == 8) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k_nope))) : "memory"); } + if (warp == 8) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_k_rope))) : "memory"); } + if (warp == 8) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_v))) : "memory"); } + + // Mbarrier init (14 pipeline groups, 0 ordered-sequence groups, 48 barriers) + // Mbarriers at smem_raw[0..384) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + mbarrier_init(smem + 112, 1); + mbarrier_init(smem + 120, 1); + mbarrier_init(smem + 128, 1); + // kv_empty: 15 barriers, init_count=1 + mbarrier_init(smem + 136, 1); + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + mbarrier_init(smem + 160, 1); + mbarrier_init(smem + 168, 1); + mbarrier_init(smem + 176, 1); + mbarrier_init(smem + 184, 1); + mbarrier_init(smem + 192, 1); + mbarrier_init(smem + 200, 1); + mbarrier_init(smem + 208, 1); + mbarrier_init(smem + 216, 1); + mbarrier_init(smem + 224, 1); + mbarrier_init(smem + 232, 1); + mbarrier_init(smem + 240, 1); + mbarrier_init(smem + 248, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 256, 1); + mbarrier_init(smem + 264, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 272, 256); + mbarrier_init(smem + 280, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 288, 256); + mbarrier_init(smem + 296, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 304, 1); + mbarrier_init(smem + 312, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 320, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 328, 128); + mbarrier_init(smem + 336, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 344, 128); + mbarrier_init(smem + 352, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 360, 1); + // tmem_scrubbed: 1 barriers, init_count=256 + mbarrier_init(smem + 368, 256); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 376, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 384); + if (warp == 0) { + int _tmem_hold = smem + 384; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + #pragma unroll 1 + for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { + int batch_idx = work_idx / (unsigned int)num_split; + int split_idx = work_idx % (unsigned int)num_split; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int tiles_per_split = (total_kv_tiles + num_split - 1) / num_split; + int start_tile = split_idx * tiles_per_split; + int end_tile_raw = start_tile + tiles_per_split; + int end_tile = ((end_tile_raw < total_kv_tiles) ? end_tile_raw : total_kv_tiles); + int num_kv_tiles = end_tile - start_tile; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + #pragma unroll 1 + for (int tile = 0; tile < num_kv_tiles; tile++) { + int pipeline_tile = softmax_tile_cursor + tile; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(n_half * 64) + (unsigned int)(tmem_row_base << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile + tile; + int tail_valid = seqlen_kv_b - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -MLA_INF; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float _max_0 = max_noftz(_tmem_load_0_max, row_max_val); + smem_softmax_exchange[exchange_idx] = _max_0; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(_max_0, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float no_correction = (((_max_1 - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - _max_1); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = _max_1; + float safe_max = ((_max_1 == -MLA_INF) ? 0.0f : _max_1); + float max_scaled = safe_max * softmax_scale_log2; + #pragma unroll + for (int i = 0; i < 64; i++) { + float _exp2_1 = approx_exp2(_tmem_load_0[i] * softmax_scale_log2 - max_scaled); + _tmem_load_0[i] = _exp2_1; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_3); + softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_3); + float _tmem_load_0_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + _tmem_load_0_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (phase) * 8, p_cor_empty_phase); + float meta[4]; + meta[0] = row_sum_val; + meta[1] = row_max_val; + meta[2] = acc_scale; + meta[3] = no_correction; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + tmem_st_x4_f32(meta_addr, meta); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t _tmem_load_0_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(_tmem_load_0_bf16[pv_off])), "r"(__as_u32(_tmem_load_0_bf16[pv_off + 1])), "r"(__as_u32(_tmem_load_0_bf16[pv_off + 2])), "r"(__as_u32(_tmem_load_0_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { + int batch_idx_1 = work_idx_1 / (unsigned int)num_split; + int split_idx_1 = work_idx_1 % (unsigned int)num_split; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int tiles_per_split_1 = (total_kv_tiles_1 + num_split - 1) / num_split; + int start_tile_1 = split_idx_1 * tiles_per_split_1; + int end_tile_raw_1 = start_tile_1 + tiles_per_split_1; + int end_tile_1 = ((end_tile_raw_1 < total_kv_tiles_1) ? end_tile_raw_1 : total_kv_tiles_1); + int num_kv_tiles_1 = end_tile_1 - start_tile_1; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile_1 = correction_tile_cursor + tile_1; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], meta_addr_1); + final_sum_val = _tmem_load_1[0]; + final_max_val = _tmem_load_1[1]; + float acc_scale_1 = _tmem_load_1[2]; + float no_correction_1 = _tmem_load_1[3]; + if (tile_1 > 0) { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + if (tile_1 > 0) { + int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); + int skip_correction = _vote_0; + if (skip_correction == 0) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + #pragma unroll + for (int acc_stage = 0; acc_stage < 2; acc_stage++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_2)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + int stat_offset = (batch_idx_1 * 128 + head_idx) * num_split + split_idx_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + *(reinterpret_cast(partial_lse + stat_offset) + (0)) = lse_log2; + int o_offset = stat_offset * 512; + float output_scale = bmm2_scale; + if (enable_sink != 0 && num_split == 1) { + float _exp2_2 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - lse_log2); + float sink_exp = _exp2_2; + float _rcp_1 = approx_rcp(1.0f + sink_exp); + output_scale = output_scale * _rcp_1; + } + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_1 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_3[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_3[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_3[j + 0], _tmem_load_3[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_3[j + 2], _tmem_load_3[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_3[j + 4], _tmem_load_3[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_3[j + 6], _tmem_load_3[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_1; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + { + __threadfence(); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int wg2_dummy = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + unsigned int mma_kv_stage = 0; + int mma_tile_cursor = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + #pragma unroll 1 + for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { + int batch_idx_2 = work_idx_2 / (unsigned int)num_split; + int split_idx_2 = work_idx_2 % (unsigned int)num_split; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int tiles_per_split_2 = (total_kv_tiles_2 + num_split - 1) / num_split; + int start_tile_2 = split_idx_2 * tiles_per_split_2; + int end_tile_raw_2 = start_tile_2 + tiles_per_split_2; + int end_tile_2 = ((end_tile_raw_2 < total_kv_tiles_2) ? end_tile_raw_2 : total_kv_tiles_2); + int num_kv_tiles_2 = end_tile_2 - start_tile_2; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_2 = mma_tile_cursor + tile_2; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n = 0; n < 8; n++) { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n == 0) ? 0 : 1))); + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_1 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + if (tile_2 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int p_stage_1 = prev_phase * 4 + ps_1; + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int output_col = 128 + acc_stage_2 * 128; + int _mma_a_lo_2 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_2 = ((((smem_v_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_2 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_2 = 0; ps_2 < 4; ps_2++) { + int p_stage_last = last_phase * 4 + ps_2; + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int output_col_d = 128 + acc_stage_3 * 128; + int _mma_a_lo_3 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_3 = ((((smem_v_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_2 == 0) ? first_pv : 0)) ? 0 : 1))); + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_2; + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy_1 = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + unsigned int load_kv_stage = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { + int batch_idx_3 = work_idx_3 / (unsigned int)num_split; + int split_idx_3 = work_idx_3 % (unsigned int)num_split; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int tiles_per_split_3 = (total_kv_tiles_3 + num_split - 1) / num_split; + int start_tile_3 = split_idx_3 * tiles_per_split_3; + int end_tile_raw_3 = start_tile_3 + tiles_per_split_3; + int end_tile_3 = ((end_tile_raw_3 < total_kv_tiles_3) ? end_tile_raw_3 : total_kv_tiles_3); + int num_kv_tiles_3 = end_tile_3 - start_tile_3; + int pt_base = batch_idx_3 * max_pages_per_seq; + int q_row_global = batch_idx_3 * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, (&tmap_q), 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile_3 + 2 * cta_rank; + int pg_k00 = page_table[k0_page_base]; + int pg_k01 = page_table[k0_page_base + 1]; + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + int dst = smem_kv_addr + load_kv_stage * 8192; + tma_4d_gmem2smem_cta2(dst, (&tmap_k_nope), 0, 0, n_1, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, (&tmap_k_nope), 0, 0, n_1, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + int dst_r = smem_kv_addr + load_kv_stage * 8192; + tma_4d_gmem2smem_cta2(dst_r, (&tmap_k_rope), 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, (&tmap_k_rope), 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { + int abs_tile_k = start_tile_3 + tile_3; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = page_table[k_page_base]; + int pg_k1 = page_table[k_page_base + 1]; + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + int dst_k = smem_kv_addr + load_kv_stage * 8192; + tma_4d_gmem2smem_cta2(dst_k, (&tmap_k_nope), 0, 0, n_2, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, (&tmap_k_nope), 0, 0, n_2, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + int dst_kr = smem_kv_addr + load_kv_stage * 8192; + tma_4d_gmem2smem_cta2(dst_kr, (&tmap_k_rope), 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, (&tmap_k_rope), 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + int abs_tile_v = start_tile_3 + tile_3 - 1; + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps_3]; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < 2; acc_stage_4++) { + int v_group = acc_stage_4 * 4 + cta_rank * 2; + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + int dst_v = smem_v_addr + load_kv_stage * 8192; + tma_4d_gmem2smem_cta2(dst_v, (&tmap_v), 0, 0, v_group, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_v + 4096, (&tmap_v), 0, 0, v_group + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + int abs_tile_last = start_tile_3 + num_kv_tiles_3 - 1; + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_4]; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < 2; acc_stage_5++) { + int v_group_last = acc_stage_5 * 4 + cta_rank * 2; + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + int dst_vl = smem_v_addr + load_kv_stage * 8192; + tma_4d_gmem2smem_cta2(dst_vl, (&tmap_v), 0, 0, v_group_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vl + 4096, (&tmap_v), 0, 0, v_group_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + } + // ---- Role: empty0 ---- + if (warp == 10) { + { // empty0_main + const int wg2_dummy_2 = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + #pragma unroll 1 + for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + // ---- Role: empty1 ---- + if (warp == 11) { + { // empty1_main + const int wg2_dummy_3 = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + #pragma unroll 1 + for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_v32_reduce_1db619cce9.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_v32_reduce_1db619cce9.cu new file mode 100644 index 00000000000..07dbbef74fb --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_v32_reduce_1db619cce9.cu @@ -0,0 +1,740 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_SPLIT_WEIGHTS_OFF 0 +#define SMEM_SPLIT_WEIGHTS_STAGE_BYTES 128 +#define SMEM_SPLIT_WEIGHTS_STRIDE 128 +#define SMEM_MERGED_LSE_OFF 128 +#define SMEM_MERGED_LSE_STAGE_BYTES 4 +#define SMEM_MERGED_LSE_STRIDE 4 +#define SMEM_TOTAL 256 +#define THREADS 128 +#define USE_PDL 1 + +#include + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_mla_decode_v32_reduce(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, float* __restrict__ sinks, int num_q_heads, int num_split, float bmm2_scale, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + float* split_weights = reinterpret_cast(smem_raw + 0); + const int split_weights_addr = smem + 0; + float* merged_lse = reinterpret_cast(smem_raw + 128); + const int merged_lse_addr = smem + 128; + + // === Task calls (dependency order) === + const int wg_dummy = 0; + int batch_idx = blockIdx.x; + int head_idx = blockIdx.y; + int stat_base = (batch_idx * num_q_heads + head_idx) * num_split; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + if (warp == 0) { + float local_m[1]; + float local_w[1]; + float thread_max = -MLA_INF; + #pragma unroll + for (int i = 0; i < 1; i++) { + int split_idx = lane + (unsigned int)(i * 32); + local_m[i] = -MLA_INF; + local_w[i] = 0.0f; + if (split_idx < num_split) { + local_m[i] = partial_lse[stat_base + split_idx]; + float _max_0 = max_noftz(thread_max, local_m[i]); + thread_max = _max_0; + } + } + float _warp_reduce_0 = thread_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + float global_max = _warp_reduce_0; + float thread_sum = 0.0f; + #pragma unroll + for (int i_1 = 0; i_1 < 1; i_1++) { + int split_idx_1 = lane + (unsigned int)(i_1 * 32); + if (split_idx_1 < num_split) { + float _exp2_0 = approx_exp2(local_m[i_1] - global_max); + local_w[i_1] = ((local_m[i_1] == -MLA_INF) ? 0.0f : _exp2_0); + thread_sum = thread_sum + local_w[i_1]; + } + } + float _warp_reduce_1 = thread_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + float global_sum = _warp_reduce_1; + float _rcp_0 = approx_rcp(global_sum); + float inv_sum = ((global_sum > 0.0f) ? _rcp_0 : 0.0f); + if (lane == 0) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(global_sum)); + merged_lse[0] = global_max + _log2_0; + } + #pragma unroll + for (int i_2 = 0; i_2 < 1; i_2++) { + int split_idx_2 = lane + (unsigned int)(i_2 * 32); + if (split_idx_2 < num_split) { + split_weights[split_idx_2] = local_w[i_2] * inv_sum; + } + } + } + __syncthreads(); + float output_scale = bmm2_scale; + if (enable_sink != 0) { + float _exp2_1 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - merged_lse[0]); + float sink_exp = _exp2_1; + float _rcp_1 = approx_rcp(1.0f + sink_exp); + output_scale = output_scale * _rcp_1; + } + int po_head_base = stat_base * 512; + int o_head_base = (batch_idx * num_q_heads + head_idx) * 512; + float acc[4]; + int d_base = tid * 4; + #pragma unroll + for (int e = 0; e < 4; e++) { + acc[e] = 0.0f; + } + #pragma unroll 4 + for (int s = 0; s < num_split; s++) { + float split_weight = split_weights[s]; + float _vec_load_0[4]; + { + uint2 _vld_0; + _vld_0 = *reinterpret_cast(partial_O + (po_head_base + s * 512 + d_base) + 0); + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&_vec_load_0[0 + _pair * 2])[0]), "=f"((&_vec_load_0[0 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + #pragma unroll + for (int e_1 = 0; e_1 < 4; e_1++) { + acc[e_1] = acc[e_1] + split_weight * _vec_load_0[e_1]; + } + } + { + const float2 _prescale2_1 = {output_scale, output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 2; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&acc[0])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + acc[0 + _ps] *= output_scale; + #endif + uint2 _pk2; + __nv_bfloat162* _pk = reinterpret_cast<__nv_bfloat162*>(&_pk2); + _pk[0] = __floats2bfloat162_rn(acc[0 + 0], acc[0 + 1]); + _pk[1] = __floats2bfloat162_rn(acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (o_head_base + d_base)))[0]) = _pk2; + } + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_value_split_paged_bf16_5c36fe8228.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_value_split_paged_bf16_5c36fe8228.cu new file mode 100644 index 00000000000..fc15cbe7a76 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_value_split_paged_bf16_5c36fe8228.cu @@ -0,0 +1,2185 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 74752 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 74752 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 74752 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 197632 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 230400 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 1024 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 1024 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 231424 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_TOTAL 231936 +#define THREADS 384 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384) __cluster_dims__(2,1,1) void +kernel_mla_decode_value_split_paged_bf16(MlaTensorMap const* tmap_q, MlaTensorMap const* tmap_k_nope, MlaTensorMap const* tmap_k_rope, MlaTensorMap const* tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (tid == 0) { + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k_nope)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k_rope)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_v)) : "memory"); + } + __syncthreads(); + + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_kv_addr = smem + 74752; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_vhalf_addr = smem + 74752; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_vquarter_addr = smem + 74752; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 197632); + const int smem_p_addr = smem + 197632; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 230400); + const int smem_softmax_exchange_addr = smem + 230400; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 231424); + const int smem_epilogue_exchange_addr = smem + 231424; + + // Mbarrier init (14 pipeline groups, 0 ordered-sequence groups, 48 barriers) + // Mbarriers at smem_raw[0..384) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + mbarrier_init(smem + 112, 1); + mbarrier_init(smem + 120, 1); + mbarrier_init(smem + 128, 1); + // kv_empty: 15 barriers, init_count=1 + mbarrier_init(smem + 136, 1); + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + mbarrier_init(smem + 160, 1); + mbarrier_init(smem + 168, 1); + mbarrier_init(smem + 176, 1); + mbarrier_init(smem + 184, 1); + mbarrier_init(smem + 192, 1); + mbarrier_init(smem + 200, 1); + mbarrier_init(smem + 208, 1); + mbarrier_init(smem + 216, 1); + mbarrier_init(smem + 224, 1); + mbarrier_init(smem + 232, 1); + mbarrier_init(smem + 240, 1); + mbarrier_init(smem + 248, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 256, 1); + mbarrier_init(smem + 264, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 272, 256); + mbarrier_init(smem + 280, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 288, 256); + mbarrier_init(smem + 296, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 304, 1); + mbarrier_init(smem + 312, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 320, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 328, 128); + mbarrier_init(smem + 336, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 344, 128); + mbarrier_init(smem + 352, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 360, 1); + // tmem_scrubbed: 1 barriers, init_count=256 + mbarrier_init(smem + 368, 256); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 376, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 384); + if (warp == 0) { + int _tmem_hold = smem + 384; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + #pragma unroll 1 + for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { + int batch_idx = work_idx; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + #pragma unroll 1 + for (int tile = 0; tile < num_kv_tiles; tile++) { + int pipeline_tile = softmax_tile_cursor + tile; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile + tile; + bool score_tile_full = 0; + bool all_tiles_are_complete_k = seqlen_kv_b % 128 == 0; + int tile_end_k = (abs_tile + 1) * 128; + bool is_full_tile_k = tile_end_k < seqlen_kv_b; + score_tile_full = all_tiles_are_complete_k || is_full_tile_k; + if (!score_tile_full) { + int tail_valid = seqlen_kv_b - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + } + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[phase * 128 + exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[phase * 128 + exchange_idx ^ 64]); + new_max = _max_1; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 scale_pair = _f2_0; + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 neg_max_pair = _f2_1; + #pragma unroll + for (int i = 0; i < 64; i += 2) { + float2 _f2_2 = make_float2(sv[i], sv[i + 1]); + float2 score_pair = _f2_2; + float2 affine_pair = fma_f32x2_rn_ftz(score_pair, scale_pair, neg_max_pair); + float _exp2_1 = approx_exp2(affine_pair.x); + sv[i] = _exp2_1; + float _exp2_2 = approx_exp2(affine_pair.y); + sv[i + 1] = _exp2_2; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (phase) * 8, p_cor_empty_phase); + float meta[4]; + meta[0] = row_sum_val; + meta[1] = row_max_val; + meta[2] = acc_scale; + meta[3] = no_correction; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + tmem_st_x4_f32(meta_addr, meta); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { + int batch_idx_1 = work_idx_1; + int value_split = 0; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + if (num_kv_tiles_1 > 0) { + int first_phase = correction_tile_cursor & 1; + int first_wait_phase = correction_tile_cursor >> 1 & 1; + mbarrier_wait(stats_addr + (first_phase) * 8, first_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int first_meta_addr = taddr + 384 + (unsigned int)(first_phase * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], first_meta_addr); + final_sum_val = _tmem_load_0[0]; + final_max_val = _tmem_load_0[1]; + mbarrier_arrive(stats_empty_addr + (first_phase) * 8); + } + #pragma unroll 1 + for (int tile_1 = 1; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile_1 = correction_tile_cursor + tile_1; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], meta_addr_1); + final_sum_val = _tmem_load_1[0]; + final_max_val = _tmem_load_1[1]; + float acc_scale_1 = _tmem_load_1[2]; + float no_correction_1 = _tmem_load_1[3]; + { + } + { + int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); + int skip_correction = _vote_0; + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + if (skip_correction == 0) { + { + { + #pragma unroll + for (int acc_stage = 0; acc_stage < 2; acc_stage++) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage * 128) + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_3[16]; + tmem_ld_x16(&_tmem_load_3[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_3)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_3)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_3); + } + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_1 * 128 + head_idx) * 512; + float output_scale = bmm2_scale; + if (enable_sink != 0) { + float _exp2_3 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - lse_log2); + float sink_exp = _exp2_3; + float _rcp_1 = approx_rcp(1.0f + sink_exp); + output_scale = output_scale * _rcp_1; + } + { + { + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = acc_stage_1 * 2 * 2 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + { + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + const float2 _scale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast(_tmem_load_8)[_ls] = mul_f32x2_rn_noftz(reinterpret_cast(_tmem_load_8)[_ls], _scale2_1); + unsigned int pair_packed[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_8[_lp*2 + 0], _tmem_load_8[_lp*2+1 + 0])); + pair_packed[_lp] = *(uint32_t*)&_bf2; + } + int pair_parity = lane & 1; + unsigned int pair_mask = (unsigned int)(-pair_parity); + int pair_neighbor = lane ^ 1; + int pair_even_base = gmem_base - pair_parity * 512 + pair_parity * 8; + #pragma unroll + for (int pair_columns = 0; pair_columns < 2; pair_columns++) { + unsigned int pair_even_words[4]; + unsigned int pair_odd_words[4]; + #pragma unroll + for (int pair_word = 0; pair_word < 4; pair_word++) { + unsigned int _shfl_0 = __shfl_sync(0xFFFFFFFF, pair_packed[pair_columns * 8 + 4 + pair_word] & ~pair_mask | pair_packed[pair_columns * 8 + pair_word] & pair_mask, pair_neighbor); + pair_even_words[pair_word] = pair_packed[pair_columns * 8 + pair_word] & ~pair_mask | _shfl_0 & pair_mask; + pair_odd_words[pair_word] = _shfl_0 & ~pair_mask | pair_packed[pair_columns * 8 + 4 + pair_word] & pair_mask; + } + reinterpret_cast(O + (pair_even_base + pair_columns * 16))[0] = reinterpret_cast(pair_even_words)[0]; + reinterpret_cast(O + (pair_even_base + 512 + pair_columns * 16))[0] = reinterpret_cast(pair_odd_words)[0]; + } + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_1; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_k_nope)) : "memory"); + { + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_k_rope)) : "memory"); + } + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_v)) : "memory"); + const int wg2_dummy = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + #pragma unroll 1 + for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { + int batch_idx_2 = work_idx_2; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_2 = mma_tile_cursor + tile_2; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n = 0; n < 8; n++) { + { + uint32_t _mbar_token_6 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_6; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (n) * 512; + int _mma_b_lo_0 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((n == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + { + uint32_t _mbar_token_7 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_7; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_2 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_2 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + if (tile_2 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int p_stage_1 = prev_phase * 4 + ps_1; + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < ((1) ? 2 : 1); acc_stage_2++) { + { + uint32_t _mbar_token_8 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_8; + int output_col = 128 + acc_stage_2 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_4 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_2 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_2 = 0; ps_2 < 4; ps_2++) { + int p_stage_last = last_phase * 4 + ps_2; + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < ((1) ? 2 : 1); acc_stage_3++) { + { + uint32_t _mbar_token_9 = mbarrier_try_wait(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase); + mma_kv_token = _mbar_token_9; + int output_col_d = 128 + acc_stage_3 * 128; + mbarrier_wait_token(kv_full_addr + (mma_kv_stage) * 8, mma_kv_phase, mma_kv_token); + int _mma_a_lo_7 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_7 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_2 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; mma_kv_phase ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + { + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy_1 = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { + int batch_idx_3 = work_idx_3; + int value_split_1 = 0; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + int pt_base = batch_idx_3 * max_pages_per_seq; + int q_row_global = batch_idx_3 * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, tmap_q, 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile_3 + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst = 0; + { + uint32_t _mbar_token_0 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_0; + dst = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst, tmap_k_nope, 0, 0, n_1, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, tmap_k_nope, 0, 0, n_1, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_r = 0; + { + uint32_t _mbar_token_1 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_1; + dst_r = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_r, tmap_k_rope, 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, tmap_k_rope, 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { + int abs_tile_k = start_tile_3 + tile_3; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + int dst_k = 0; + { + uint32_t _mbar_token_2 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_2; + dst_k = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_k, tmap_k_nope, 0, 0, n_2, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, tmap_k_nope, 0, 0, n_2, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + int dst_kr = 0; + { + uint32_t _mbar_token_3 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_3; + dst_kr = smem_kv_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_kr, tmap_k_rope, 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, tmap_k_rope, 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + int abs_tile_v = start_tile_3 + tile_3 - 1; + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps_3]; + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < ((1) ? 2 : 1); acc_stage_4++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + uint32_t _mbar_token_4 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_4; + v_group_half = acc_stage_4 * 4 + cta_rank * 2; + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf, tmap_v, 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, tmap_v, 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + int abs_tile_last = start_tile_3 + num_kv_tiles_3 - 1; + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_4]; + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < ((1) ? 2 : 1); acc_stage_5++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + uint32_t _mbar_token_5 = mbarrier_try_wait(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase); + load_kv_token = _mbar_token_5; + v_group_half_last = acc_stage_5 * 4 + cta_rank * 2; + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + mbarrier_wait_token(kv_empty_addr + (load_kv_stage) * 8, load_kv_phase, load_kv_token); + } + { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(dst_vhalf_last, tmap_v, 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, tmap_v, 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + } + // ---- Role: empty0 ---- + if (warp == 10) { + { // empty0_main + const int wg2_dummy_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + // ---- Role: empty1 ---- + if (warp == 11) { + { // empty1_main + const int wg2_dummy_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_value_split_paged_bf16_ee41d3ceea.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_value_split_paged_bf16_ee41d3ceea.cu new file mode 100644 index 00000000000..7762b0b396e --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_value_split_paged_bf16_ee41d3ceea.cu @@ -0,0 +1,2130 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 74752 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 74752 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 74752 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 197632 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 230400 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 230912 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_TOTAL 231424 +#define THREADS 384 +#define V_QUARTER 1 +#define UNSPLIT 0 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384) __cluster_dims__(2,1,1) void +kernel_mla_decode_value_split_paged_bf16(MlaTensorMap const* tmap_q, MlaTensorMap const* tmap_k_nope, MlaTensorMap const* tmap_k_rope, MlaTensorMap const* tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (tid == 0) { + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k_nope)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k_rope)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_v)) : "memory"); + } + __syncthreads(); + + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_kv_addr = smem + 74752; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_vhalf_addr = smem + 74752; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_vquarter_addr = smem + 74752; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 197632); + const int smem_p_addr = smem + 197632; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 230400); + const int smem_softmax_exchange_addr = smem + 230400; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 230912); + const int smem_epilogue_exchange_addr = smem + 230912; + + // Mbarrier init (14 pipeline groups, 0 ordered-sequence groups, 48 barriers) + // Mbarriers at smem_raw[0..384) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + mbarrier_init(smem + 112, 1); + mbarrier_init(smem + 120, 1); + mbarrier_init(smem + 128, 1); + // kv_empty: 15 barriers, init_count=1 + mbarrier_init(smem + 136, 1); + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + mbarrier_init(smem + 160, 1); + mbarrier_init(smem + 168, 1); + mbarrier_init(smem + 176, 1); + mbarrier_init(smem + 184, 1); + mbarrier_init(smem + 192, 1); + mbarrier_init(smem + 200, 1); + mbarrier_init(smem + 208, 1); + mbarrier_init(smem + 216, 1); + mbarrier_init(smem + 224, 1); + mbarrier_init(smem + 232, 1); + mbarrier_init(smem + 240, 1); + mbarrier_init(smem + 248, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 256, 1); + mbarrier_init(smem + 264, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 272, 256); + mbarrier_init(smem + 280, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 288, 256); + mbarrier_init(smem + 296, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 304, 1); + mbarrier_init(smem + 312, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 320, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 328, 128); + mbarrier_init(smem + 336, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 344, 128); + mbarrier_init(smem + 352, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 360, 1); + // tmem_scrubbed: 1 barriers, init_count=256 + mbarrier_init(smem + 368, 256); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 376, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 384); + if (warp == 0) { + int _tmem_hold = smem + 384; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + #pragma unroll 1 + for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { + int batch_idx = work_idx / (unsigned int)value_split_count; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + #pragma unroll 1 + for (int tile = 0; tile < num_kv_tiles; tile++) { + int pipeline_tile = softmax_tile_cursor + tile; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile + tile; + bool score_tile_full = 0; + int tail_valid = seqlen_kv_b - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + #pragma unroll + for (int i = 0; i < 64; i++) { + float _exp2_1 = approx_exp2(sv[i] * softmax_scale_log2 - max_scaled); + sv[i] = _exp2_1; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (phase) * 8, p_cor_empty_phase); + float meta[4]; + meta[0] = row_sum_val; + meta[1] = row_max_val; + meta[2] = acc_scale; + meta[3] = no_correction; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + tmem_st_x4_f32(meta_addr, meta); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { + int batch_idx_1 = work_idx_1 / (unsigned int)value_split_count; + int value_split = work_idx_1 % (unsigned int)value_split_count; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile_1 = correction_tile_cursor + tile_1; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], meta_addr_1); + final_sum_val = _tmem_load_0[0]; + final_max_val = _tmem_load_0[1]; + float acc_scale_1 = _tmem_load_0[2]; + float no_correction_1 = _tmem_load_0[3]; + if (tile_1 > 0) { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + if (tile_1 > 0) { + int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + #pragma unroll + for (int vs_local = 0; vs_local < 1; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_1[16]; + tmem_ld_x16(&_tmem_load_1[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_1)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_1)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_1); + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_1 * 128 + head_idx) * 512; + float output_scale = bmm2_scale; + if (enable_sink != 0) { + float _exp2_2 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - lse_log2); + float sink_exp = _exp2_2; + float _rcp_1 = approx_rcp(1.0f + sink_exp); + output_scale = output_scale * _rcp_1; + } + { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 1; vs_local_1++) { + int logical_vs = value_split * 2 + n_half_1 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_4[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]), "=f"(_tmem_load_4[2]), "=f"(_tmem_load_4[3]), "=f"(_tmem_load_4[4]), "=f"(_tmem_load_4[5]), "=f"(_tmem_load_4[6]), "=f"(_tmem_load_4[7]), "=f"(_tmem_load_4[8]), "=f"(_tmem_load_4[9]), "=f"(_tmem_load_4[10]), "=f"(_tmem_load_4[11]), "=f"(_tmem_load_4[12]), "=f"(_tmem_load_4[13]), "=f"(_tmem_load_4[14]), "=f"(_tmem_load_4[15]), "=f"(_tmem_load_4[16]), "=f"(_tmem_load_4[17]), "=f"(_tmem_load_4[18]), "=f"(_tmem_load_4[19]), "=f"(_tmem_load_4[20]), "=f"(_tmem_load_4[21]), "=f"(_tmem_load_4[22]), "=f"(_tmem_load_4[23]), "=f"(_tmem_load_4[24]), "=f"(_tmem_load_4[25]), "=f"(_tmem_load_4[26]), "=f"(_tmem_load_4[27]), "=f"(_tmem_load_4[28]), "=f"(_tmem_load_4[29]), "=f"(_tmem_load_4[30]), "=f"(_tmem_load_4[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_4[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_4[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_4[j + 0], _tmem_load_4[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_4[j + 2], _tmem_load_4[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_4[j + 4], _tmem_load_4[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_4[j + 6], _tmem_load_4[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_1; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_k_nope)) : "memory"); + { + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_k_rope)) : "memory"); + } + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_v)) : "memory"); + const int wg2_dummy = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + #pragma unroll 1 + for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { + int batch_idx_2 = work_idx_2 / (unsigned int)value_split_count; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_2 = mma_tile_cursor + tile_2; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n = 0; n < 8; n++) { + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (n) * 512; + int _mma_b_lo_1 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(((n == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int _mma_a_lo_3 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_3 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + if (tile_2 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int p_stage_1 = prev_phase * 4 + ps_1; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((UNSPLIT) ? 2 : 1); acc_stage++) { + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int output_col = 128 + acc_stage * 128; + { + int _mma_a_lo_5 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_5 = ((((smem_vquarter_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136381584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_2 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_2 = 0; ps_2 < 4; ps_2++) { + int p_stage_last = last_phase * 4 + ps_2; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((UNSPLIT) ? 2 : 1); acc_stage_1++) { + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int output_col_d = 128 + acc_stage_1 * 128; + { + int _mma_a_lo_8 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_8 = ((((smem_vquarter_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136381584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_8), "r"(_mma_b_lo_8), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_2 == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + { + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy_1 = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { + int batch_idx_3 = work_idx_3 / (unsigned int)value_split_count; + int value_split_1 = work_idx_3 % (unsigned int)value_split_count; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + int pt_base = batch_idx_3 * max_pages_per_seq; + int q_row_global = batch_idx_3 * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, tmap_q, 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile_3 + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst, tmap_k_nope, 0, 0, n_1, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, tmap_k_nope, 0, 0, n_1, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + int dst_r = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_r = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_r, tmap_k_rope, 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, tmap_k_rope, 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { + int abs_tile_k = start_tile_3 + tile_3; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + int dst_k = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_k = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_k, tmap_k_nope, 0, 0, n_2, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, tmap_k_nope, 0, 0, n_2, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + int dst_kr = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_kr = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_kr, tmap_k_rope, 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, tmap_k_rope, 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + int abs_tile_v = start_tile_3 + tile_3 - 1; + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps_3]; + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < ((UNSPLIT) ? 2 : 1); acc_stage_2++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + { + int v_group_quarter = value_split_1 * 2 + cta_rank; + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(8192)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(smem_vquarter_addr + load_kv_stage * 8192, tmap_v, 0, 0, v_group_quarter, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + } + int abs_tile_last = start_tile_3 + num_kv_tiles_3 - 1; + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_4]; + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < ((UNSPLIT) ? 2 : 1); acc_stage_3++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + { + int v_group_quarter_last = value_split_1 * 2 + cta_rank; + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(8192)) : "memory"); + } + } + if (elect_sync()) { + tma_4d_gmem2smem_cta2(smem_vquarter_addr + load_kv_stage * 8192, tmap_v, 0, 0, v_group_quarter_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + } + // ---- Role: empty0 ---- + if (warp == 10) { + { // empty0_main + const int wg2_dummy_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + // ---- Role: empty1 ---- + if (warp == 11) { + { // empty1_main + const int wg2_dummy_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_value_split_paged_bf16_fa7b59f3a7.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_value_split_paged_bf16_fa7b59f3a7.cu new file mode 100644 index 00000000000..409849d1a77 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_mla_decode_value_split_paged_bf16_fa7b59f3a7.cu @@ -0,0 +1,2146 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_KV_PIPE_STAGES 15 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 74752 +#define SMEM_SMEM_KV_STAGE_BYTES 8192 +#define SMEM_SMEM_KV_STRIDE 8192 +#define SMEM_SMEM_VHALF_OFF 74752 +#define SMEM_SMEM_VHALF_STAGE_BYTES 8192 +#define SMEM_SMEM_VHALF_STRIDE 8192 +#define SMEM_SMEM_VQUARTER_OFF 74752 +#define SMEM_SMEM_VQUARTER_STAGE_BYTES 4096 +#define SMEM_SMEM_VQUARTER_STRIDE 8192 +#define SMEM_SMEM_P_OFF 197632 +#define SMEM_SMEM_P_STAGE_BYTES 4096 +#define SMEM_SMEM_P_STRIDE 4096 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 230400 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 230912 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_TOTAL 231424 +#define THREADS 384 +#define V_QUARTER 0 +#define UNSPLIT 0 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(384) __cluster_dims__(2,1,1) void +kernel_mla_decode_value_split_paged_bf16(MlaTensorMap const* tmap_q, MlaTensorMap const* tmap_k_nope, MlaTensorMap const* tmap_k_rope, MlaTensorMap const* tmap_v, __nv_bfloat16* __restrict__ O, int* __restrict__ seq_lens_kv, int* __restrict__ page_table, float* __restrict__ sinks, float softmax_scale_log2, float bmm2_scale, int total_work_items, int value_split_count, int max_pages_per_seq, int enable_sink) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 136) + #define s_full_addr (mbar_base + 256) + #define s_empty_addr (mbar_base + 272) + #define p_full_addr (mbar_base + 288) + #define p_empty_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 320) + #define stats_addr (mbar_base + 328) + #define stats_empty_addr (mbar_base + 344) + #define o_full_addr (mbar_base + 360) + #define tmem_scrubbed_addr (mbar_base + 368) + #define tmem_dealloc_peer_addr (mbar_base + 376) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (tid == 0) { + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k_nope)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k_rope)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_v)) : "memory"); + } + __syncthreads(); + + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_kv_addr = smem + 74752; + __nv_bfloat16* smem_vhalf = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_vhalf_addr = smem + 74752; + __nv_bfloat16* smem_vquarter = reinterpret_cast<__nv_bfloat16*>(smem_raw + 74752); + const int smem_vquarter_addr = smem + 74752; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 197632); + const int smem_p_addr = smem + 197632; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 230400); + const int smem_softmax_exchange_addr = smem + 230400; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 230912); + const int smem_epilogue_exchange_addr = smem + 230912; + + // Mbarrier init (14 pipeline groups, 0 ordered-sequence groups, 48 barriers) + // Mbarriers at smem_raw[0..384) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 15 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + mbarrier_init(smem + 112, 1); + mbarrier_init(smem + 120, 1); + mbarrier_init(smem + 128, 1); + // kv_empty: 15 barriers, init_count=1 + mbarrier_init(smem + 136, 1); + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + mbarrier_init(smem + 160, 1); + mbarrier_init(smem + 168, 1); + mbarrier_init(smem + 176, 1); + mbarrier_init(smem + 184, 1); + mbarrier_init(smem + 192, 1); + mbarrier_init(smem + 200, 1); + mbarrier_init(smem + 208, 1); + mbarrier_init(smem + 216, 1); + mbarrier_init(smem + 224, 1); + mbarrier_init(smem + 232, 1); + mbarrier_init(smem + 240, 1); + mbarrier_init(smem + 248, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 256, 1); + mbarrier_init(smem + 264, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 272, 256); + mbarrier_init(smem + 280, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 288, 256); + mbarrier_init(smem + 296, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 304, 1); + mbarrier_init(smem + 312, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 320, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 328, 128); + mbarrier_init(smem + 336, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 344, 128); + mbarrier_init(smem + 352, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 360, 1); + // tmem_scrubbed: 1 barriers, init_count=256 + mbarrier_init(smem + 368, 256); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 376, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 384); + if (warp == 0) { + int _tmem_hold = smem + 384; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int n_half = warp % 4 / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + const int seed_row_base = warp % 4 * 32; + int softmax_tile_cursor = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_col = 0; seed_col < 512; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(seed_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_scrubbed_addr); + int seed_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_scrubbed_addr), "r"(seed_peer_rank) : "memory"); + #pragma unroll 1 + for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { + int batch_idx = work_idx / (unsigned int)value_split_count; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int total_kv_tiles = (seqlen_kv_b + 128 - 1) / 128; + int start_tile = 0; + int num_kv_tiles = total_kv_tiles; + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + #pragma unroll 1 + for (int tile = 0; tile < num_kv_tiles; tile++) { + int pipeline_tile = softmax_tile_cursor + tile; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 64 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_row_base << 16); + float sv[64]; + float tile_max = -MLA_INF; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_base + 32)); + int abs_tile = start_tile + tile; + bool score_tile_full = 0; + int tail_valid = seqlen_kv_b - abs_tile * 128 - n_half * 64; + uint32_t _slice_lo_mask_0; + { + int _lim_0 = tail_valid; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = tail_valid - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + float new_max = tile_max; + float2 _reg_reduce_max2_2 = {-MLA_INF, -MLA_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv[32], _reg_reduce_max2_2); + float sv_max = row_max_reduce(_reg_reduce_max2_2); + new_max = sv_max; + float _max_0 = max_noftz(new_max, row_max_val); + new_max = _max_0; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float _max_1 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + new_max = _max_1; + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + #pragma unroll + for (int i = 0; i < 64; i++) { + float _exp2_1 = approx_exp2(sv[i] * softmax_scale_log2 - max_scaled); + sv[i] = _exp2_1; + } + float2 _reg_reduce_sum2_3 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_3); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_3); + float sv_sum = _reg_reduce_sum2_3.x + _reg_reduce_sum2_3.y; + row_sum_val = row_sum_val * acc_scale + sv_sum; + int p_cor_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (phase) * 8, p_cor_empty_phase); + float meta[4]; + meta[0] = row_sum_val; + meta[1] = row_max_val; + meta[2] = acc_scale; + meta[3] = no_correction; + int meta_addr = taddr + 384 + (unsigned int)(phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(tmem_row_base << 16); + tmem_st_x4_f32(meta_addr, meta); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + int p_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + const int p_row = my_row; + uint32_t sv_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(sv[_lp*2 + 0], sv[_lp*2+1 + 0])); + sv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int ps_local = 0; ps_local < 2; ps_local++) { + int ps = n_half * 2 + ps_local; + int p_stage = phase * 4 + ps; + int p_base_smem = smem_p_addr + (unsigned int)(p_stage * 4096); + #pragma unroll + for (int vec = 0; vec < 4; vec++) { + int pv_off = ps_local * 16 + vec * 4; + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (p_row * 64 + vec * 16 ^ (p_row * 64 + vec * 16 >> 7 & 3) << 4))), "r"(__as_u32(sv_bf16[pv_off])), "r"(__as_u32(sv_bf16[pv_off + 1])), "r"(__as_u32(sv_bf16[pv_off + 2])), "r"(__as_u32(sv_bf16[pv_off + 3])) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 208;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { + int batch_idx_1 = work_idx_1 / (unsigned int)value_split_count; + int value_split = work_idx_1 % (unsigned int)value_split_count; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int total_kv_tiles_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int start_tile_1 = 0; + int num_kv_tiles_1 = total_kv_tiles_1; + float final_sum_val = 0.0f; + float final_max_val = -MLA_INF; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile_1 = correction_tile_cursor + tile_1; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], meta_addr_1); + final_sum_val = _tmem_load_0[0]; + final_max_val = _tmem_load_0[1]; + float acc_scale_1 = _tmem_load_0[2]; + float no_correction_1 = _tmem_load_0[3]; + if (tile_1 > 0) { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + if (tile_1 > 0) { + int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); + int skip_correction = _vote_0; + if (skip_correction == 0) { + { + { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int o_base = taddr + 128 + (unsigned int)(vs_local * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_3[16]; + tmem_ld_x16(&_tmem_load_3[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + reinterpret_cast(_tmem_load_3)[_ls] = mul_f32x2_rn_ftz(reinterpret_cast(_tmem_load_3)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_3); + } + } + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int epilogue_exchange_idx = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[epilogue_exchange_idx] = final_sum_val; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + final_sum_val = final_sum_val + smem_epilogue_exchange[epilogue_exchange_idx ^ 64]; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum_val); + float inv_sum = _rcp_0; + int head_idx = cta_rank * 64 + my_row_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum_val)); + float lse_log2 = final_max_val * softmax_scale_log2 + _log2_0; + int o_offset = (batch_idx_1 * 128 + head_idx) * 512; + float output_scale = bmm2_scale; + if (enable_sink != 0) { + float _exp2_2 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - lse_log2); + float sink_exp = _exp2_2; + float _rcp_1 = approx_rcp(1.0f + sink_exp); + output_scale = output_scale * _rcp_1; + } + { + { + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int logical_vs = value_split * 4 + n_half_1 * 2 + vs_local_1; + int o_base_epi = taddr + 128 + (unsigned int)(vs_local_1 * 64) + (unsigned int)corr_row; + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_8[0]), "=f"(_tmem_load_8[1]), "=f"(_tmem_load_8[2]), "=f"(_tmem_load_8[3]), "=f"(_tmem_load_8[4]), "=f"(_tmem_load_8[5]), "=f"(_tmem_load_8[6]), "=f"(_tmem_load_8[7]), "=f"(_tmem_load_8[8]), "=f"(_tmem_load_8[9]), "=f"(_tmem_load_8[10]), "=f"(_tmem_load_8[11]), "=f"(_tmem_load_8[12]), "=f"(_tmem_load_8[13]), "=f"(_tmem_load_8[14]), "=f"(_tmem_load_8[15]), "=f"(_tmem_load_8[16]), "=f"(_tmem_load_8[17]), "=f"(_tmem_load_8[18]), "=f"(_tmem_load_8[19]), "=f"(_tmem_load_8[20]), "=f"(_tmem_load_8[21]), "=f"(_tmem_load_8[22]), "=f"(_tmem_load_8[23]), "=f"(_tmem_load_8[24]), "=f"(_tmem_load_8[25]), "=f"(_tmem_load_8[26]), "=f"(_tmem_load_8[27]), "=f"(_tmem_load_8[28]), "=f"(_tmem_load_8[29]), "=f"(_tmem_load_8[30]), "=f"(_tmem_load_8[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + logical_vs * 64 + c_1; + #pragma unroll + for (int j = 0; j < 32; j += 8) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_8[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_8[j + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_8[j + 0], _tmem_load_8[j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_8[j + 2], _tmem_load_8[j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_8[j + 4], _tmem_load_8[j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_8[j + 6], _tmem_load_8[j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_1; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_k_nope)) : "memory"); + { + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_k_rope)) : "memory"); + } + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_v)) : "memory"); + const int wg2_dummy = 0; + unsigned int mma_kv_stage = 0; + unsigned int mma_kv_phase = 0; + unsigned int mma_kv_token = 0; + int mma_tile_cursor = 0; + unsigned int _phase_tmem_scrubbed_0 = 0; + mbarrier_wait(tmem_scrubbed_addr, _phase_tmem_scrubbed_0); + _phase_tmem_scrubbed_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_o_empty_0 = 1; + #pragma unroll 1 + for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { + int batch_idx_2 = work_idx_2 / (unsigned int)value_split_count; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int total_kv_tiles_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int start_tile_2 = 0; + int num_kv_tiles_2 = total_kv_tiles_2; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_2 = mma_tile_cursor + tile_2; + int phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + int score_col = ((phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int n = 0; n < 8; n++) { + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (n) * 512; + int _mma_b_lo_1 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(((n == 0) ? 0 : 1))); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int _mma_a_lo_3 = (((smem_q_addr) >> 4) & 0x3FFF) + (8) * 512; + int _mma_b_lo_3 = (((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + if (tile_2 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_1 = 0; ps_1 < 4; ps_1++) { + int p_stage_1 = prev_phase * 4 + ps_1; + #pragma unroll + for (int acc_stage = 0; acc_stage < ((UNSPLIT) ? 2 : 1); acc_stage++) { + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int output_col = 128 + acc_stage * 128; + { + int _mma_a_lo_6 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_1) * 256; + int _mma_b_lo_6 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_2 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_2 = 0; ps_2 < 4; ps_2++) { + int p_stage_last = last_phase * 4 + ps_2; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < ((UNSPLIT) ? 2 : 1); acc_stage_1++) { + { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + int output_col_d = 128 + acc_stage_1 * 128; + { + int _mma_a_lo_9 = (((smem_p_addr) >> 4) & 0x3FFF) + (p_stage_last) * 256; + int _mma_b_lo_9 = ((((smem_vhalf_addr) >> 4) & 0x3FFF) | 0x1000000) + (mma_kv_stage) * 512; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_9), "r"(_mma_b_lo_9), "r"((tmem_tmem_scratch + (output_col_d))), "r"(((((ps_2 == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(kv_empty_addr + (mma_kv_stage) * 8, (uint16_t)(3)); + { + mma_kv_stage += 1; + if (mma_kv_stage == 15) { mma_kv_stage = 0; _phase_kv_full ^= 1; } + } + } + } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + { + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_2; + } + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = mma_tile_cursor + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_warp ---- + if (warp == 9) { + { // load_warp_main + const int wg2_dummy_1 = 0; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + unsigned int load_kv_token = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_kv_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { + int batch_idx_3 = work_idx_3 / (unsigned int)value_split_count; + int value_split_1 = work_idx_3 % (unsigned int)value_split_count; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int total_kv_tiles_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int start_tile_3 = 0; + int num_kv_tiles_3 = total_kv_tiles_3; + int pt_base = batch_idx_3 * max_pages_per_seq; + int q_row_global = batch_idx_3 * 128 + cta_rank * 64; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(147456)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int s = 0; s < 9; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_3d_gmem2smem_cta2(q_dst, tmap_q, 0, q_row_global, s, ((q_full_addr) & 0xFEFFFFFF)); + } + } + int k0_page_base = pt_base + 4 * start_tile_3 + 2 * cta_rank; + int pg_k00 = 0; + int pg_k01 = 0; + { + pg_k00 = page_table[k0_page_base]; + pg_k01 = page_table[k0_page_base + 1]; + } + #pragma unroll + for (int n_1 = 0; n_1 < 8; n_1++) { + int dst = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst, tmap_k_nope, 0, 0, n_1, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst + 4096, tmap_k_nope, 0, 0, n_1, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + int dst_r = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_r = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_r, tmap_k_rope, 0, 0, 8, pg_k00, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_r + 4096, tmap_k_rope, 0, 0, 8, pg_k01, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { + int abs_tile_k = start_tile_3 + tile_3; + int k_page_base = pt_base + 4 * abs_tile_k + 2 * cta_rank; + int pg_k0 = 0; + int pg_k1 = 0; + { + pg_k0 = page_table[k_page_base]; + pg_k1 = page_table[k_page_base + 1]; + } + #pragma unroll + for (int n_2 = 0; n_2 < 8; n_2++) { + int dst_k = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_k = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_k, tmap_k_nope, 0, 0, n_2, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_k + 4096, tmap_k_nope, 0, 0, n_2, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + int dst_kr = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_kr = smem_kv_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_kr, tmap_k_rope, 0, 0, 8, pg_k0, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_kr + 4096, tmap_k_rope, 0, 0, 8, pg_k1, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + int abs_tile_v = start_tile_3 + tile_3 - 1; + #pragma unroll + for (int ps_3 = 0; ps_3 < 4; ps_3++) { + int pg_v = page_table[pt_base + 4 * abs_tile_v + ps_3]; + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < ((UNSPLIT) ? 2 : 1); acc_stage_2++) { + int v_group_half = 0; + int dst_vhalf = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + { + { + v_group_half = value_split_1 * 4 + cta_rank * 2; + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_vhalf = smem_vhalf_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_vhalf, tmap_v, 0, 0, v_group_half, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf + 4096, tmap_v, 0, 0, v_group_half + 1, pg_v, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + } + int abs_tile_last = start_tile_3 + num_kv_tiles_3 - 1; + #pragma unroll + for (int ps_4 = 0; ps_4 < 4; ps_4++) { + int pg_v_last = page_table[pt_base + 4 * abs_tile_last + ps_4]; + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < ((UNSPLIT) ? 2 : 1); acc_stage_3++) { + int v_group_half_last = 0; + int dst_vhalf_last = 0; + { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + } + { + { + v_group_half_last = value_split_1 * 4 + cta_rank * 2; + } + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(16384)) : "memory"); + } + } + if (elect_sync()) { + { + dst_vhalf_last = smem_vhalf_addr + load_kv_stage * 8192; + } + tma_4d_gmem2smem_cta2(dst_vhalf_last, tmap_v, 0, 0, v_group_half_last, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(dst_vhalf_last + 4096, tmap_v, 0, 0, v_group_half_last + 1, pg_v_last, ((kv_full_addr + (load_kv_stage) * 8) & 0xFEFFFFFF)); + } + } + { + load_kv_stage += 1; + if (load_kv_stage == 15) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + } + // ---- Role: empty0 ---- + if (warp == 10) { + { // empty0_main + const int wg2_dummy_2 = 0; + #pragma unroll 1 + for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + // ---- Role: empty1 ---- + if (warp == 11) { + { // empty1_main + const int wg2_dummy_3 = 0; + #pragma unroll 1 + for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_page64_keeps_gmem_separate_reduce_cea01a87e9.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_page64_keeps_gmem_separate_reduce_cea01a87e9.cu new file mode 100644 index 00000000000..667f682d4b4 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_page64_keeps_gmem_separate_reduce_cea01a87e9.cu @@ -0,0 +1,704 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 512 +#define USE_PDL 1 + +#include + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +extern "C" { + +__global__ __launch_bounds__(512) void +kernel_page64_keeps_gmem_separate_reduce(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_stats, __nv_bfloat16* __restrict__ O, float* __restrict__ LSE, int num_split, int num_reduce_ctas, int num_queries, float bmm2_scale) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + const int role_token = 0; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + int reduction_cta = blockIdx.x; + int work_item = blockIdx.z; + int partition = blockIdx.y; + int head_group = partition; + int v_partition = 0; + int slices_per_cta = (8 + num_reduce_ctas - 1) / num_reduce_ctas; + int slice_begin = reduction_cta * slices_per_cta; + int slice_end_raw = slice_begin + slices_per_cta; + int slice_end = ((slice_end_raw < 8) ? slice_end_raw : 8); + #pragma unroll 1 + for (int slice_idx = slice_begin; slice_idx < slice_end; slice_idx++) { + int vector_idx = tid; + int row_in_group = slice_idx * 8 + vector_idx / 64; + int dim = v_partition * 512 + vector_idx % 64 * 8; + int head = head_group * 64 + row_in_group; + float merged[8]; + #pragma unroll + for (int item = 0; item < 8; item++) { + merged[item] = 0.0f; + } + float merged_max = -MLA_INF; + float merged_sum = 0.0f; + #pragma unroll 4 + for (int split = 0; split < num_split; split++) { + int stat_base = ((work_item * num_split + split) * 128 + head) * 2; + float local_max = partial_stats[stat_base]; + float local_sum = partial_stats[stat_base + 1]; + float _max_0 = max_noftz(merged_max, local_max); + float next_max = _max_0; + float _exp2_0 = approx_exp2(merged_max - next_max); + float prior_scale = ((merged_max != -MLA_INF) ? _exp2_0 : 0.0f); + float _exp2_1 = approx_exp2(local_max - next_max); + float local_scale = ((local_max != -MLA_INF) ? _exp2_1 : 0.0f); + int partial_base = ((work_item * num_split + split) * 128 + head) * 512 + dim; + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(partial_O + partial_base + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&_vec_load_0[0 + _blk * 8 + _pair * 2])[0]), "=f"((&_vec_load_0[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + #pragma unroll + for (int item_1 = 0; item_1 < 8; item_1++) { + merged[item_1] = merged[item_1] * prior_scale + _vec_load_0[item_1] * local_scale; + } + merged_sum = merged_sum * prior_scale + local_sum * local_scale; + merged_max = next_max; + } + float _rcp_0 = approx_rcp(merged_sum); + float inv_sum = ((merged_sum > 0.0f) ? _rcp_0 : 0.0f); + int output_base = (work_item * 128 + head) * 512 + dim; + { + const float2 _prescale2_1 = {inv_sum * bmm2_scale, inv_sum * bmm2_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&merged[0])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + merged[0 + _ps] *= inv_sum * bmm2_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(merged[0 + 0], merged[0 + 1]); + _pk[1] = __floats2bfloat162_rn(merged[0 + 2], merged[0 + 3]); + _pk[2] = __floats2bfloat162_rn(merged[0 + 4], merged[0 + 5]); + _pk[3] = __floats2bfloat162_rn(merged[0 + 6], merged[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + output_base))[0]) = *reinterpret_cast(&_pk[0]); + } + if (dim == 0) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(merged_sum)); + LSE[work_item * 128 + head] = ((merged_sum > 0.0f) ? merged_max + _log2_0 : -MLA_INF); + } + } + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_trtllm_mla_decode_fp8_6fe1e0de09.cu b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_trtllm_mla_decode_fp8_6fe1e0de09.cu new file mode 100644 index 00000000000..107091f73fa --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/device/sm_103a_152/cake_trtllm_mla_decode_fp8_6fe1e0de09.cu @@ -0,0 +1,2284 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "MLA requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) MlaTensorMap { uint64_t opaque[16]; }; +struct __align__(64) MlaTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(MlaTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(MlaTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) MlaTensorMapPack { MlaTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(MlaTensorMap) >= alignof(CUtensorMap), "MlaTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define MLA_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_K_PIPE_STAGES 2 +#define NUM_V_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 40960 +#define SMEM_SMEM_Q_STRIDE 40960 +#define SMEM_SMEM_K_OFF 41984 +#define SMEM_SMEM_K_STAGE_BYTES 40960 +#define SMEM_SMEM_K_STRIDE 40960 +#define SMEM_SMEM_V_OFF 123904 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_SMEM_STATS_MAX_OFF 189440 +#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 +#define SMEM_SMEM_STATS_MAX_STRIDE 1024 +#define SMEM_SMEM_STATS_SUM_OFF 190464 +#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_SUM_STRIDE 512 +#define SMEM_TOTAL 190976 +#define THREADS 384 +#define USE_PDL 1 +#define SM103_SCORE_LD_RED_MAX 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ uint32_t mbarrier_test_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.test_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_test_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +__device__ __forceinline__ void mbarrier_wait_token_test(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_TEST:\n\t" + "mbarrier.test_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_TEST;\n\t" + "bra.uni LAB_WAIT_TEST;\n\t" + "DONE_TEST:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_test_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_TEST_CLUSTER:\n\t" + "mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_TEST_CLUSTER;\n\t" + "bra.uni LAB_WAIT_TEST_CLUSTER;\n\t" + "DONE_TEST_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x32_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x32.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16," + " %17, %18, %19, %20, %21, %22, %23, %24," + " %25, %26, %27, %28, %29, %30, %31, %32};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15]), + "f"(src[16]), "f"(src[17]), "f"(src[18]), "f"(src[19]), + "f"(src[20]), "f"(src[21]), "f"(src[22]), "f"(src[23]), + "f"(src[24]), "f"(src[25]), "f"(src[26]), "f"(src[27]), + "f"(src[28]), "f"(src[29]), "f"(src[30]), "f"(src[31])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" + :: "r"(addr), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), + "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); +} + +extern "C" { + +__global__ __launch_bounds__(384, 1) __cluster_dims__(2,1,1) void +kernel_trtllm_mla_decode_fp8(MlaTensorMap const* tmap_q, MlaTensorMap const* tmap_k, MlaTensorMap const* tmap_v, __nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_stats, __nv_bfloat16* __restrict__ O, float* __restrict__ LSE, int* __restrict__ seq_lens_kv, int* __restrict__ work_batch_indices, int* __restrict__ page_table, float softmax_scale_log2, float bmm2_scale, int num_split, int num_queries, int max_pages_per_seq) +{ + const int tid = threadIdx.x; + const uint32_t warp = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane = static_cast(tid) & 31u; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define k_full_addr (mbar_base + 16) + #define k_empty_addr (mbar_base + 32) + #define v_full_addr (mbar_base + 48) + #define v_empty_addr (mbar_base + 64) + #define s_full_addr (mbar_base + 80) + #define s_empty_addr (mbar_base + 96) + #define p_full_addr (mbar_base + 112) + #define p_empty_addr (mbar_base + 128) + #define o_empty_addr (mbar_base + 144) + #define stats_addr (mbar_base + 152) + #define stats_empty_addr (mbar_base + 168) + #define sum_ready_addr (mbar_base + 184) + #define o_full_addr (mbar_base + 192) + #define partial_results_ready_addr (mbar_base + 200) + #define q_pair_ready_addr (mbar_base + 208) + #define tmem_dealloc_peer_addr (mbar_base + 216) + #define tmem_epoch_ready_addr (mbar_base + 224) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (tid == 0) { + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_k)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_v)) : "memory"); + } + __syncthreads(); + + + // Kernel setup ops + uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + uint8_t* smem_k = reinterpret_cast(smem_raw + 41984); + const int smem_k_addr = smem + 41984; + uint8_t* smem_v = reinterpret_cast(smem_raw + 123904); + const int smem_v_addr = smem + 123904; + float* smem_stats_max = reinterpret_cast(smem_raw + 189440); + const int smem_stats_max_addr = smem + 189440; + float* smem_stats_sum = reinterpret_cast(smem_raw + 190464); + const int smem_stats_sum_addr = smem + 190464; + + // Mbarrier init (19 pipeline groups, 0 ordered-sequence groups, 29 barriers) + // Mbarriers at smem_raw[0..232) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'k_pipe' --- + // k_full: 2 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + // k_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + // --- pipeline 'v_pipe' --- + // v_full: 2 barriers, init_count=1 + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + // v_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 96, 256); + mbarrier_init(smem + 104, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 112, 256); + mbarrier_init(smem + 120, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 128, 1); + mbarrier_init(smem + 136, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 144, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 152, 128); + mbarrier_init(smem + 160, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 168, 128); + mbarrier_init(smem + 176, 128); + // sum_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 184, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 192, 1); + // partial_results_ready: 1 barriers, init_count=256 + mbarrier_init(smem + 200, 256); + // q_pair_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 208, 64); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 216, 32); + // tmem_epoch_ready: 1 barriers, init_count=256 + mbarrier_init(smem + 224, 256); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 232); + if (warp == 0) { + int _tmem_hold = smem + 232; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 112;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 200;"); + { // softmax_wg_main + uint32_t _mbar_token_0 = mbarrier_test_wait(q_full_addr, 0); + const int wg_dummy_inc = 0; + const int tmem_row_base = warp % 2 * 32; + const int score_half = warp / 2; + const int my_row = (unsigned int)tmem_row_base + lane; + int split_idx = blockIdx.x / 2; + int batch_idx = blockIdx.y; + int seqlen_kv_b = seq_lens_kv[batch_idx]; + int num_kv_tiles_total = (seqlen_kv_b + 128 - 1) / 128; + int blocks_per_split = (num_kv_tiles_total + num_split - 1) / num_split; + int my_start_block = split_idx * blocks_per_split; + int my_end_block_raw = my_start_block + blocks_per_split; + int my_end_block = ((my_end_block_raw < num_kv_tiles_total) ? my_end_block_raw : num_kv_tiles_total); + int my_n_tiles = my_end_block - my_start_block; + mbarrier_wait_token_test(q_full_addr, 0, _mbar_token_0); + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + float row_max_val = -MLA_INF; + float row_sum_val = 0.0f; + #pragma unroll 1 + for (int tile = 0; tile < my_n_tiles; tile++) { + int phase = tile & 1; + int s_wait_phase = tile >> 1 & 1; + float sv[64]; + float local_max = -MLA_INF; + { + uint32_t _mbar_token_1 = mbarrier_test_wait(s_full_addr + (phase) * 8, s_wait_phase); + int s_off = ((phase != 0) ? 128 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(score_half * 64) + (unsigned int)(tmem_row_base << 16); + int abs_tile = my_start_block + tile; + int tail_valid = seqlen_kv_b - abs_tile * 128; + int local_tail_raw = tail_valid - score_half * 64; + int local_tail_hi = ((local_tail_raw > 64) ? 64 : local_tail_raw); + int local_tail = ((local_tail_hi < 0) ? 0 : local_tail_hi); + if (local_tail < 64) { + mbarrier_wait_token(s_full_addr + (phase) * 8, s_wait_phase, _mbar_token_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]), "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]), "=f"(local_max) + : "r"(s_base)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + uint32_t _slice_lo_mask_0; + { + int _lim_0 = local_tail; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv[0] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv[1] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv[2] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv[3] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv[4] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv[5] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv[6] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv[7] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv[8] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv[9] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv[10] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv[11] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv[12] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv[13] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv[14] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv[15] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv[16] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv[17] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv[18] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv[19] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv[20] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv[21] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv[22] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv[23] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv[24] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv[25] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv[26] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv[27] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv[28] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv[29] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv[30] = -MLA_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv[31] = -MLA_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = local_tail - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv[32] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv[33] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv[34] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv[35] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv[36] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv[37] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv[38] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv[39] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv[40] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv[41] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv[42] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv[43] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv[44] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv[45] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv[46] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv[47] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv[48] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv[49] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv[50] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv[51] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv[52] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv[53] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv[54] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv[55] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv[56] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv[57] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv[58] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv[59] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv[60] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv[61] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv[62] = -MLA_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv[63] = -MLA_INF; + local_max = -MLA_INF; + #pragma unroll + for (int i = 0; i < 64; i++) { + float _max_0 = max_noftz(local_max, sv[i]); + local_max = _max_0; + } + } else { + mbarrier_wait_token_test(s_full_addr + (phase) * 8, s_wait_phase, _mbar_token_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]), "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]), "=f"(local_max) + : "r"(s_base)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + } + } + int stats_empty_phase = tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (phase) * 8, stats_empty_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int pair_max_idx = phase * 128 + score_half * 64 + my_row; + smem_stats_max[pair_max_idx] = local_max; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (warp % 2 == 0) { + asm volatile("barrier.sync 1, 64;" ::: "memory"); + } else { + asm volatile("barrier.sync 2, 64;" ::: "memory"); + } + asm volatile("tcgen05.fence::after_thread_sync;"); + int partner_max_idx = phase * 128 + (1 - score_half) * 64 + my_row; + float partner_max = smem_stats_max[partner_max_idx]; + float _max_2 = max_noftz(local_max, partner_max); + float new_max = _max_2; + float _max_3 = max_noftz(new_max, row_max_val); + new_max = _max_3; + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -MLA_INF) ? exp_delta : 1.0f); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (warp >= 2) { + smem_stats_max[phase * 128 + my_row] = acc_scale; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + row_max_val = new_max; + float safe_max = ((new_max == -MLA_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + #pragma unroll + for (int i_1 = 0; i_1 < 64; i_1++) { + float _exp2_1 = approx_exp2(sv[i_1] * softmax_scale_log2 - max_scaled); + sv[i_1] = _exp2_1; + } + float local_sum = 0.0f; + #pragma unroll + for (int i_2 = 0; i_2 < 64; i_2++) { + local_sum = local_sum + sv[i_2]; + } + row_sum_val = row_sum_val * acc_scale + local_sum; + int p_empty_phase = tile >> 1 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (phase) * 8, p_empty_phase); + int p_off = ((phase != 0) ? 224 : 96); + int p_base = taddr + (unsigned int)p_off + (unsigned int)(score_half * 16) + (unsigned int)(tmem_row_base << 16); + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(sv[0 + _j * 4]), "f"(sv[0 + _j * 4 + 1]), + "f"(sv[0 + _j * 4 + 2]), "f"(sv[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; + } + tmem_st_x8_u32(p_base, _pv_packed); + } + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(sv[32 + _j * 4]), "f"(sv[32 + _j * 4 + 1]), + "f"(sv[32 + _j * 4 + 2]), "f"(sv[32 + _j * 4 + 3])); + _pv_packed[_j] = _pk; + } + tmem_st_x8_u32(p_base + 8, _pv_packed); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(phase * 8)), "r"(0) : "memory"); + } + int final_sum_slot = my_row * 2 + score_half; + smem_stats_sum[final_sum_slot] = row_sum_val; + if (score_half == 0) { + smem_stats_max[64 + my_row] = row_max_val; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(sum_ready_addr); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(partial_results_ready_addr); + unsigned int _phase_partial_results_ready_0 = 0; + mbarrier_wait(partial_results_ready_addr, _phase_partial_results_ready_0); + _phase_partial_results_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // correction_wg_main + const int wg_dummy_inc_1 = 0; + const int tmem_row_base_1 = warp % 2 * 32; + const int o_tmem_row_base = warp % 4 * 32; + const int n_half = warp % 4 / 2; + const int my_row_1 = (unsigned int)tmem_row_base_1 + lane; + const int corr_row = o_tmem_row_base << 16; + float scrub_zero[128]; + #pragma unroll + for (int scrub_i = 0; scrub_i < 128; scrub_i++) { + scrub_zero[scrub_i] = 0.0f; + } + const int scrub_row = (unsigned int)o_tmem_row_base + lane; + #pragma unroll + for (int scrub_col = 0; scrub_col < 512; scrub_col += 128) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x128.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, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %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};" + :: "r"(taddr + (unsigned int)scrub_col + (unsigned int)(scrub_row << 16)), "f"(scrub_zero[0]), "f"(scrub_zero[1]), "f"(scrub_zero[2]), "f"(scrub_zero[3]), "f"(scrub_zero[4]), "f"(scrub_zero[5]), "f"(scrub_zero[6]), "f"(scrub_zero[7]), "f"(scrub_zero[8]), "f"(scrub_zero[9]), "f"(scrub_zero[10]), "f"(scrub_zero[11]), "f"(scrub_zero[12]), "f"(scrub_zero[13]), "f"(scrub_zero[14]), "f"(scrub_zero[15]), "f"(scrub_zero[16]), "f"(scrub_zero[17]), "f"(scrub_zero[18]), "f"(scrub_zero[19]), "f"(scrub_zero[20]), "f"(scrub_zero[21]), "f"(scrub_zero[22]), "f"(scrub_zero[23]), "f"(scrub_zero[24]), "f"(scrub_zero[25]), "f"(scrub_zero[26]), "f"(scrub_zero[27]), "f"(scrub_zero[28]), "f"(scrub_zero[29]), "f"(scrub_zero[30]), "f"(scrub_zero[31]), "f"(scrub_zero[32]), "f"(scrub_zero[33]), "f"(scrub_zero[34]), "f"(scrub_zero[35]), "f"(scrub_zero[36]), "f"(scrub_zero[37]), "f"(scrub_zero[38]), "f"(scrub_zero[39]), "f"(scrub_zero[40]), "f"(scrub_zero[41]), "f"(scrub_zero[42]), "f"(scrub_zero[43]), "f"(scrub_zero[44]), "f"(scrub_zero[45]), "f"(scrub_zero[46]), "f"(scrub_zero[47]), "f"(scrub_zero[48]), "f"(scrub_zero[49]), "f"(scrub_zero[50]), "f"(scrub_zero[51]), "f"(scrub_zero[52]), "f"(scrub_zero[53]), "f"(scrub_zero[54]), "f"(scrub_zero[55]), "f"(scrub_zero[56]), "f"(scrub_zero[57]), "f"(scrub_zero[58]), "f"(scrub_zero[59]), "f"(scrub_zero[60]), "f"(scrub_zero[61]), "f"(scrub_zero[62]), "f"(scrub_zero[63]), "f"(scrub_zero[64]), "f"(scrub_zero[65]), "f"(scrub_zero[66]), "f"(scrub_zero[67]), "f"(scrub_zero[68]), "f"(scrub_zero[69]), "f"(scrub_zero[70]), "f"(scrub_zero[71]), "f"(scrub_zero[72]), "f"(scrub_zero[73]), "f"(scrub_zero[74]), "f"(scrub_zero[75]), "f"(scrub_zero[76]), "f"(scrub_zero[77]), "f"(scrub_zero[78]), "f"(scrub_zero[79]), "f"(scrub_zero[80]), "f"(scrub_zero[81]), "f"(scrub_zero[82]), "f"(scrub_zero[83]), "f"(scrub_zero[84]), "f"(scrub_zero[85]), "f"(scrub_zero[86]), "f"(scrub_zero[87]), "f"(scrub_zero[88]), "f"(scrub_zero[89]), "f"(scrub_zero[90]), "f"(scrub_zero[91]), "f"(scrub_zero[92]), "f"(scrub_zero[93]), "f"(scrub_zero[94]), "f"(scrub_zero[95]), "f"(scrub_zero[96]), "f"(scrub_zero[97]), "f"(scrub_zero[98]), "f"(scrub_zero[99]), "f"(scrub_zero[100]), "f"(scrub_zero[101]), "f"(scrub_zero[102]), "f"(scrub_zero[103]), "f"(scrub_zero[104]), "f"(scrub_zero[105]), "f"(scrub_zero[106]), "f"(scrub_zero[107]), "f"(scrub_zero[108]), "f"(scrub_zero[109]), "f"(scrub_zero[110]), "f"(scrub_zero[111]), "f"(scrub_zero[112]), "f"(scrub_zero[113]), "f"(scrub_zero[114]), "f"(scrub_zero[115]), "f"(scrub_zero[116]), "f"(scrub_zero[117]), "f"(scrub_zero[118]), "f"(scrub_zero[119]), "f"(scrub_zero[120]), "f"(scrub_zero[121]), "f"(scrub_zero[122]), "f"(scrub_zero[123]), "f"(scrub_zero[124]), "f"(scrub_zero[125]), "f"(scrub_zero[126]), "f"(scrub_zero[127])); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_epoch_ready_addr); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_epoch_ready_addr), "r"(cta_rank ^ 1) : "memory"); + int split_idx_1 = blockIdx.x / 2; + int batch_idx_1 = blockIdx.y; + int seqlen_kv_b_1 = seq_lens_kv[batch_idx_1]; + int num_kv_tiles_total_1 = (seqlen_kv_b_1 + 128 - 1) / 128; + int blocks_per_split_1 = (num_kv_tiles_total_1 + num_split - 1) / num_split; + int my_start_block_1 = split_idx_1 * blocks_per_split_1; + int my_end_block_raw_1 = my_start_block_1 + blocks_per_split_1; + int my_end_block_1 = ((my_end_block_raw_1 < num_kv_tiles_total_1) ? my_end_block_raw_1 : num_kv_tiles_total_1); + int my_n_tiles_1 = my_end_block_1 - my_start_block_1; + unsigned int _phase_q_full_0 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < my_n_tiles_1; tile_1++) { + int phase_1 = tile_1 & 1; + int stats_wait_phase = tile_1 >> 1 & 1; + float acc_scale_1 = 1.0f; + { + uint32_t _mbar_token_2 = mbarrier_test_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + int stats_max_idx = phase_1 * 128 + my_row_1; + if (tile_1 > 0) { + int o_ready_phase = tile_1 - 1 & 1; + uint32_t _mbar_token_3 = mbarrier_test_wait(o_full_addr, o_ready_phase); + mbarrier_wait_token(stats_addr + (phase_1) * 8, stats_wait_phase, _mbar_token_2); + asm volatile("tcgen05.fence::after_thread_sync;"); + acc_scale_1 = smem_stats_max[stats_max_idx]; + mbarrier_wait_token(o_full_addr, o_ready_phase, _mbar_token_3); + asm volatile("tcgen05.fence::after_thread_sync;"); + } else { + mbarrier_wait_token(stats_addr + (phase_1) * 8, stats_wait_phase, _mbar_token_2); + asm volatile("tcgen05.fence::after_thread_sync;"); + acc_scale_1 = smem_stats_max[stats_max_idx]; + } + } + if (tile_1 > 0) { + int _vote_0 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); + int any_rescale = _vote_0; + if (any_rescale != 0) { + #pragma unroll + for (int vs_local = 0; vs_local < 2; vs_local++) { + int vs = vs_local; + int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 128; c += 32) { + float _tmem_load_0[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) + : "r"(o_base + c)); + #pragma unroll + for (int j = 0; j < 32; j++) { + _tmem_load_0[j] = _tmem_load_0[j] * acc_scale_1; + } + tmem_st_x32_f32(o_base + c, _tmem_load_0); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + if (tile_1 > 0) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (phase_1) * 8); + } + { + int o_final_phase = ((my_n_tiles_1 > 0) ? my_n_tiles_1 - 1 & 1 : 0); + mbarrier_wait(o_full_addr, o_final_phase); + } + unsigned int _phase_sum_ready_0 = 0; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); + _phase_sum_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int final_sum_slot_1 = my_row_1 * 2 + n_half; + float local_final_sum = smem_stats_sum[final_sum_slot_1]; + smem_stats_sum[final_sum_slot_1] = local_final_sum; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (warp % 2 == 0) { + asm volatile("barrier.sync 4, 64;" ::: "memory"); + } else { + asm volatile("barrier.sync 5, 64;" ::: "memory"); + } + asm volatile("tcgen05.fence::after_thread_sync;"); + int partner_final_sum_slot = my_row_1 * 2 + (1 - n_half); + float partner_final_sum = smem_stats_sum[partner_final_sum_slot]; + float total_sum = local_final_sum + partner_final_sum; + float final_row_max = smem_stats_max[64 + my_row_1]; + float safe_sum = ((total_sum > 0.0f) ? total_sum : 1.0f); + float _rcp_0 = approx_rcp(safe_sum); + float output_scale = ((num_split == 1) ? _rcp_0 : 1.0f); + int head_idx = cta_rank * 64 + my_row_1; + if (n_half == 0) { + int stat_off = ((batch_idx_1 * num_split + split_idx_1) * 128 + head_idx) * 2; + float row_max_scaled = final_row_max * softmax_scale_log2; + float stored_max = ((final_row_max == -MLA_INF) ? -MLA_INF : row_max_scaled); + float packed_stats[2]; + packed_stats[0] = stored_max; + packed_stats[1] = total_sum; + { + float2 _v2 = make_float2(packed_stats[0 + 0], packed_stats[0 + 1]); + *reinterpret_cast(partial_stats + stat_off + 0) = _v2; + } + if (num_split == 1) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(total_sum)); + float direct_lse = ((total_sum > 0.0f) ? stored_max + _log2_0 : -MLA_INF); + LSE[batch_idx_1 * 128 + head_idx] = direct_lse; + } + } + int po_offset = ((batch_idx_1 * num_split + split_idx_1) * 128 + head_idx) * 512 + n_half * 256; + int direct_offset = (batch_idx_1 * 128 + head_idx) * 512 + n_half * 256; + #pragma unroll + for (int vs_local_1 = 0; vs_local_1 < 2; vs_local_1++) { + int vs_1 = vs_local_1; + int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c_1 = 0; c_1 < 128; c_1 += 64) { + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]) + : "r"(o_base_epi + c_1)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]) + : "r"(o_base_epi + c_1 + 32)); + int gmem_base = po_offset + vs_1 * 128 + c_1; + { + unsigned int packed_epi[16]; + if (num_split == 1) { + const float2 _scale2_0 = {output_scale * bmm2_scale, output_scale * bmm2_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_1 + 0))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_1 + 0))[_ls], _scale2_0); + } else { + const float2 _scale2_1 = {output_scale, output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_1 + 0))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_1 + 0))[_ls], _scale2_1); + } + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_1[_lp*2 + 0], _tmem_load_1[_lp*2+1 + 0])); + packed_epi[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int j_1 = 0; j_1 < 16; j_1 += 4) { + if (num_split == 1) { + int direct_base = direct_offset + vs_1 * 128 + c_1; + reinterpret_cast(O + (direct_base + j_1 * 2))[0] = reinterpret_cast(packed_epi + j_1)[0]; + } else { + reinterpret_cast(partial_O + (gmem_base + j_1 * 2))[0] = reinterpret_cast(packed_epi + j_1)[0]; + } + } + if (num_split == 1) { + const float2 _scale2_2 = {output_scale * bmm2_scale, output_scale * bmm2_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_1 + 32))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_1 + 32))[_ls], _scale2_2); + } else { + const float2 _scale2_3 = {output_scale, output_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + reinterpret_cast((_tmem_load_1 + 32))[_ls] = mul_f32x2_rn_ftz(reinterpret_cast((_tmem_load_1 + 32))[_ls], _scale2_3); + } + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_1[_lp*2 + 32], _tmem_load_1[_lp*2+1 + 32])); + packed_epi[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int j_2 = 0; j_2 < 16; j_2 += 4) { + if (num_split == 1) { + int direct_base_1 = direct_offset + vs_1 * 128 + c_1; + reinterpret_cast(O + (direct_base_1 + 32 + j_2 * 2))[0] = reinterpret_cast(packed_epi + j_2)[0]; + } else { + reinterpret_cast(partial_O + (gmem_base + 32 + j_2 * 2))[0] = reinterpret_cast(packed_epi + j_2)[0]; + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(partial_results_ready_addr); + unsigned int _phase_partial_results_ready_0_1 = 0; + mbarrier_wait(partial_results_ready_addr, _phase_partial_results_ready_0_1); + _phase_partial_results_ready_0_1 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + if (USE_PDL && num_split > 1) { + __threadfence(); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + asm volatile("barrier.sync 13, 160;" ::: "memory"); + if (warp == 4) { + int tmem_dealloc_peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int wg3_dummy = 0; + unsigned int mma_k_stage = 0; + unsigned int mma_v_stage = 0; + int split_idx_2 = blockIdx.x / 2; + int batch_idx_2 = blockIdx.y; + int seqlen_kv_b_2 = seq_lens_kv[batch_idx_2]; + int num_kv_tiles_total_2 = (seqlen_kv_b_2 + 128 - 1) / 128; + int blocks_per_split_2 = (num_kv_tiles_total_2 + num_split - 1) / num_split; + int my_start_block_2 = split_idx_2 * blocks_per_split_2; + int my_end_block_raw_2 = my_start_block_2 + blocks_per_split_2; + int my_end_block_2 = ((my_end_block_raw_2 < num_kv_tiles_total_2) ? my_end_block_raw_2 : num_kv_tiles_total_2); + int my_n_tiles_2 = my_end_block_2 - my_start_block_2; + unsigned int _phase_tmem_epoch_ready_0 = 0; + mbarrier_wait(tmem_epoch_ready_addr, _phase_tmem_epoch_ready_0); + _phase_tmem_epoch_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0_1 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0_1); + _phase_q_full_0_1 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(q_pair_ready_addr), "r"(0) : "memory"); + unsigned int _phase_q_pair_ready_0 = 0; + if (cta_rank == 0) { + mbarrier_wait_cluster_hint(q_pair_ready_addr, _phase_q_pair_ready_0, 10000000); + _phase_q_pair_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + int first_pv = 1; + unsigned int _phase_k_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_v_full = 0; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < my_n_tiles_2; tile_2++) { + int phase_2 = tile_2 & 1; + int s_empty_phase = tile_2 >> 1 & 1 ^ 1; + int score_col = ((phase_2 != 0) ? 128 : 0); + if (cta_rank == 0) { + mbarrier_wait(s_empty_addr + (phase_2) * 8, s_empty_phase); + mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + { + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (0) * 2560; + int _mma_b_lo_0 = (((smem_k_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 2560; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(0)); + } + { + int _mma_a_lo_2 = (((smem_q_addr + 8192) >> 4) & 0x3FFF) + (0) * 2560; + int _mma_b_lo_2 = (((smem_k_addr + 8192) >> 4) & 0x3FFF) + (mma_k_stage) * 2560; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + { + int _mma_a_lo_4 = (((smem_q_addr + 16384) >> 4) & 0x3FFF) + (0) * 2560; + int _mma_b_lo_4 = (((smem_k_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 2560; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + { + int _mma_a_lo_6 = (((smem_q_addr + 24576) >> 4) & 0x3FFF) + (0) * 2560; + int _mma_b_lo_6 = (((smem_k_addr + 24576) >> 4) & 0x3FFF) + (mma_k_stage) * 2560; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + { + int _mma_a_lo_9 = (((smem_q_addr + 32768) >> 4) & 0x3FFF) + (0) * 2560; + int _mma_b_lo_9 = (((smem_k_addr + 32768) >> 4) & 0x3FFF) + (mma_k_stage) * 2560; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_9), "r"(_mma_b_lo_9), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + } + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(k_empty_addr + (mma_k_stage) * 8, (uint16_t)(3)); + } + mma_k_stage += 1; + if (mma_k_stage == 2) { mma_k_stage = 0; _phase_k_full ^= 1; } + if (tile_2 > 0) { + int prev_phase = tile_2 - 1 & 1; + int prev_tile = tile_2 - 1; + int pv_wait_phase = prev_tile >> 1 & 1; + if (cta_rank == 0) { + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int p_col = ((prev_phase != 0) ? 224 : 96); + mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_b_lo_10 = ((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 136380432;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"((tmem_tmem_scratch + (256))), "r"(_mma_b_lo_10), "r"(tmem_tmem_scratch + p_col), "r"(((first_pv) ? 0 : 1))); + int _mma_b_lo_11 = ((((smem_v_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 136380432;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"((tmem_tmem_scratch + (384))), "r"(_mma_b_lo_11), "r"(tmem_tmem_scratch + p_col), "r"(((first_pv) ? 0 : 1))); + elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); + mma_v_stage += 1; + if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } + elect_commit_cg2_multicast(p_empty_addr + (prev_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + first_pv = 0; + } + } + int last_phase = my_n_tiles_2 - 1 & 1; + int last_tile = my_n_tiles_2 - 1; + int drain_wait_phase = last_tile >> 1 & 1; + if (cta_rank == 0) { + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int p_col_last = ((last_phase != 0) ? 224 : 96); + mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_b_lo_12 = ((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 136380432;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"((tmem_tmem_scratch + (256))), "r"(_mma_b_lo_12), "r"(tmem_tmem_scratch + p_col_last), "r"(((first_pv) ? 0 : 1))); + int _mma_b_lo_13 = ((((smem_v_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 136380432;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"((tmem_tmem_scratch + (384))), "r"(_mma_b_lo_13), "r"(tmem_tmem_scratch + p_col_last), "r"(((first_pv) ? 0 : 1))); + elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); + mma_v_stage += 1; + if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } + elect_commit_cg2_multicast(p_empty_addr + (last_phase) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("barrier.arrive 13, 160;" ::: "memory"); + } + } + // ---- Role: page_offsets_warp ---- + if (warp == 9) { + { // page_offsets_warp_main + const int wg3_dummy_1 = 0; + unsigned int _phase_q_full_0_2 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0_2); + _phase_q_full_0_2 ^= 1; + } + } + // ---- Role: padding_warp ---- + if (warp == 10) { + { // padding_warp_main + const int wg3_dummy_2 = 0; + unsigned int _phase_q_full_0_3 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0_3); + _phase_q_full_0_3 ^= 1; + } + } + // ---- Role: load_warp ---- + if (warp == 11) { + { // load_warp_main + const int wg3_dummy_3 = 0; + unsigned int load_k_stage = 0; + unsigned int load_v_stage = 0; + int split_idx_3 = blockIdx.x / 2; + int batch_idx_3 = blockIdx.y; + int seqlen_kv_b_3 = seq_lens_kv[batch_idx_3]; + int num_kv_tiles_total_3 = (seqlen_kv_b_3 + 128 - 1) / 128; + int blocks_per_split_3 = (num_kv_tiles_total_3 + num_split - 1) / num_split; + int my_start_block_3 = split_idx_3 * blocks_per_split_3; + int my_end_block_raw_3 = my_start_block_3 + blocks_per_split_3; + int my_end_block_3 = ((my_end_block_raw_3 < num_kv_tiles_total_3) ? my_end_block_raw_3 : num_kv_tiles_total_3); + int my_n_tiles_3 = my_end_block_3 - my_start_block_3; + int page_batch_idx = work_batch_indices[batch_idx_3]; + int pt_base = page_batch_idx * max_pages_per_seq; + int q_row_global = batch_idx_3 * 128 + cta_rank * 64; + unsigned int _phase_q_empty_0 = 1; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 40960); + #pragma unroll + for (int s = 0; s < 5; s++) { + int q_dst = smem_q_addr + (unsigned int)(s * 8192); + tma_2d_gmem2smem(q_dst, tmap_q, s * 128, q_row_global, q_full_addr); + } + } + int pg_k0 = page_table[pt_base + 2 * my_start_block_3 + cta_rank]; + unsigned int _phase_k_empty = 1; + mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(81920)) : "memory"); + } + } + if (elect_sync()) { + int k0_dst = smem_k_addr + load_k_stage * 40960; + #pragma unroll + for (int n = 0; n < 5; n++) { + tma_2d_gmem2smem_cta2(k0_dst + n * 8192, tmap_k, n * 128, pg_k0 * 64, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF)); + } + } + load_k_stage += 1; + if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } + unsigned int _phase_v_empty = 1; + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < my_n_tiles_3; tile_3++) { + int abs_tile_k = my_start_block_3 + tile_3; + int pg_k = page_table[pt_base + 2 * abs_tile_k + cta_rank]; + mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(81920)) : "memory"); + } + } + if (elect_sync()) { + int dst_k = smem_k_addr + load_k_stage * 40960; + #pragma unroll + for (int n_1 = 0; n_1 < 5; n_1++) { + tma_2d_gmem2smem_cta2(dst_k + n_1 * 8192, tmap_k, n_1 * 128, pg_k * 64, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF)); + } + } + load_k_stage += 1; + if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } + int abs_tile_v = my_start_block_3 + (tile_3 - 1); + int pg_v0 = page_table[pt_base + 2 * abs_tile_v]; + int pg_v1 = page_table[pt_base + 2 * abs_tile_v + 1]; + mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); + } + } + if (elect_sync()) { + int v_dst = smem_v_addr + load_v_stage * 32768; + int v0_z = cta_rank; + int v1_z = 2 + cta_rank; + tma_2d_gmem2smem_cta2(v_dst, tmap_v, v0_z * 128, pg_v0 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + tma_2d_gmem2smem_cta2(v_dst + 16384, tmap_v, v1_z * 128, pg_v0 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + tma_2d_gmem2smem_cta2(v_dst + 8192, tmap_v, v0_z * 128, pg_v1 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + tma_2d_gmem2smem_cta2(v_dst + 16384 + 8192, tmap_v, v1_z * 128, pg_v1 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + } + load_v_stage += 1; + if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } + } + int abs_tile_last = my_start_block_3 + (my_n_tiles_3 - 1); + int pg_vl0 = page_table[pt_base + 2 * abs_tile_last]; + int pg_vl1 = page_table[pt_base + 2 * abs_tile_last + 1]; + mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); + } + } + if (elect_sync()) { + int v_dst_last = smem_v_addr + load_v_stage * 32768; + int v0_z_last = cta_rank; + int v1_z_last = 2 + cta_rank; + tma_2d_gmem2smem_cta2(v_dst_last, tmap_v, v0_z_last * 128, pg_vl0 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + tma_2d_gmem2smem_cta2(v_dst_last + 16384, tmap_v, v1_z_last * 128, pg_vl0 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + tma_2d_gmem2smem_cta2(v_dst_last + 8192, tmap_v, v0_z_last * 128, pg_vl1 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + tma_2d_gmem2smem_cta2(v_dst_last + 16384 + 8192, tmap_v, v1_z_last * 128, pg_vl1 * 64, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF)); + } + load_v_stage += 1; + if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_clc.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_clc.cpp new file mode 100644 index 00000000000..28d77371093 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_clc.cpp @@ -0,0 +1,2810 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_30cfe17e1f); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_3b5bd1e042); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_7aa2c0c00a); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_8587a85908); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_8ace9da7a5); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_bbfb74c484); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_c278f217ca); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_c6609aef39); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +namespace variant_mla_decode_exact_live_bf16_clc_bbfb74c484_2e80d8d24161 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_bbfb74c484::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_bbfb74c484_2e80d8d24161 + +namespace variant_mla_decode_exact_live_bf16_clc_7aa2c0c00a_eda49f6bffff { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_7aa2c0c00a::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 231424, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 231424u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 231424u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_7aa2c0c00a_eda49f6bffff + +namespace variant_mla_decode_exact_live_bf16_clc_3b5bd1e042_c01040d8bcfe { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_3b5bd1e042::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_3b5bd1e042_c01040d8bcfe + +namespace variant_mla_decode_exact_live_bf16_clc_8ace9da7a5_ba559c721d8a { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_8ace9da7a5::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_8ace9da7a5_ba559c721d8a + +namespace variant_mla_decode_exact_live_bf16_clc_c6609aef39_924cd77cca78 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_c6609aef39::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_c6609aef39_924cd77cca78 + +namespace variant_mla_decode_exact_live_bf16_clc_8587a85908_2fcbb5c2e41d { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_8587a85908::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 231424, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 231424u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 231424u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_8587a85908_2fcbb5c2e41d + +namespace variant_mla_decode_exact_live_bf16_clc_c278f217ca_bf399908d27d { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_c278f217ca::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_c278f217ca_bf399908d27d + +namespace variant_mla_decode_exact_live_bf16_clc_30cfe17e1f_1abd96571256 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_30cfe17e1f::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_30cfe17e1f_1abd96571256 + +int SelectVariant(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z, int32_t max_kv_len, int32_t last_kv_len, int32_t source_selector_eligibility) { + if (((((total_work_items == 512) && (value_split_count == 1)) && (grid_x == 2)) && (grid_z == 512))) return 0; + if (((((((total_work_items == 4096) && (value_split_count == 4)) && (max_pages_per_seq == 32)) && (grid_x == 8)) && (grid_z == 1024)) && ((((source_selector_eligibility == 1) || (source_selector_eligibility == 3)) || (source_selector_eligibility == 5)) || (source_selector_eligibility == 7)))) return 1; + if ((((((((((source_selector_eligibility == 6) || (source_selector_eligibility == 7)) && (total_work_items == 6144)) && (value_split_count == 8)) && (max_kv_len == 1024)) && (last_kv_len == 1024)) && (max_pages_per_seq == 32)) && (grid_x == 16)) && (grid_z == 768))) return 2; + if (((((total_work_items == 6144) && (value_split_count == 8)) && (grid_x == 16)) && (grid_z == 768))) return 3; + if (((((((source_selector_eligibility == 2) || (source_selector_eligibility == 3)) || (source_selector_eligibility == 6)) || (source_selector_eligibility == 7)) && (last_kv_len == max_kv_len)) && ((((((total_work_items == 128) && (value_split_count == 2)) && (max_kv_len == 8192)) && (max_pages_per_seq == 256)) && (grid_x == 4)) && (grid_z == 64)))) return 4; + if (((((((source_selector_eligibility == 2) || (source_selector_eligibility == 3)) || (source_selector_eligibility == 6)) || (source_selector_eligibility == 7)) && (last_kv_len == max_kv_len)) && ((((((total_work_items == 8192) && (value_split_count == 16)) && (max_kv_len == 1024)) && (max_pages_per_seq == 32)) && (grid_x == 32)) && (grid_z == 512)))) return 5; + if (((((((source_selector_eligibility == 2) || (source_selector_eligibility == 3)) || (source_selector_eligibility == 6)) || (source_selector_eligibility == 7)) && (last_kv_len == max_kv_len)) && ((((((total_work_items == 1024) && (value_split_count == 2)) && (max_kv_len == 4096)) && (max_pages_per_seq == 128)) && (grid_x == 4)) && (grid_z == 512)))) return 6; + if ((((total_work_items > 0) && (grid_x > 0)) && (grid_z > 0))) return 7; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_q_rows, tvm::ffi::TensorView arg_kv_pages, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sinks, float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z, int32_t max_kv_len, int32_t last_kv_len, int32_t source_selector_eligibility, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_q_rows, "q_rows"); + CheckCudaTensor(arg_kv_pages, "kv_pages"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z, max_kv_len, last_kv_len, source_selector_eligibility); + switch (v) { + case 0: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_bbfb74c484_2e80d8d24161::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 1: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_7aa2c0c00a_eda49f6bffff::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 2: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_3b5bd1e042_c01040d8bcfe::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 3: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_8ace9da7a5_ba559c721d8a::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 4: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_c6609aef39_924cd77cca78::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 5: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_8587a85908_2fcbb5c2e41d::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 6: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_c278f217ca_bf399908d27d::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 7: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_30cfe17e1f_1abd96571256::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_bf16_clc shape (softmax_scale_log2=" << softmax_scale_log2 << ", bmm2_scale=" << bmm2_scale << ", total_work_items=" << total_work_items << ", value_split_count=" << value_split_count << ", max_pages_per_seq=" << max_pages_per_seq << ", enable_sink=" << enable_sink << ", grid_x=" << grid_x << ", grid_z=" << grid_z << ", max_kv_len=" << max_kv_len << ", last_kv_len=" << last_kv_len << ", source_selector_eligibility=" << source_selector_eligibility << ")"; + } +} + +int64_t RouteOf(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z, int32_t max_kv_len, int32_t last_kv_len, int32_t source_selector_eligibility) { return SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z, max_kv_len, last_kv_len, source_selector_eligibility); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_native_split8_pdl.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_native_split8_pdl.cpp new file mode 100644 index 00000000000..a4fcf191e84 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_native_split8_pdl.cpp @@ -0,0 +1,767 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(mla_decode_v32_paged_bf16_c80238a3c3); +TVM_FFI_EMBED_CUBIN(mla_decode_v32_reduce_1db619cce9); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +namespace variant_mla_decode_v32_paged_bf16_c80238a3c3_1fef86a30b84 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_nope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_nope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_nope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_nope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_nope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_nope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_nope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_nope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_rope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_rope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_rope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_rope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_rope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_rope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_rope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_rope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k_nope, TensorView arg_tmap_k_rope, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_num_split, int64_t arg_total_work_items, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k_nope, "tmap_k_nope"); + CheckDtype(arg_tmap_k_nope, "tmap_k_nope", 4, 16, 1); + CheckContiguous(arg_tmap_k_nope, "tmap_k_nope"); + CheckCudaTensor(arg_tmap_k_rope, "tmap_k_rope"); + CheckDtype(arg_tmap_k_rope, "tmap_k_rope", 4, 16, 1); + CheckContiguous(arg_tmap_k_rope, "tmap_k_rope"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) + << "scalar 'num_split' value " << arg_num_split + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k_nope, arg_tmap_q, "tmap_k_nope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_k_rope, arg_tmap_q, "tmap_k_rope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k_nope = EncodeTma_tmap_k_nope(arg_tmap_k_nope); + CUtensorMap p_tmap_k_rope = EncodeTma_tmap_k_rope(arg_tmap_k_rope); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_partial_lse = arg_partial_lse.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_num_split = (int32_t)arg_num_split; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k_nope, &p_tmap_k_rope, &p_tmap_v, &p_O, &p_partial_lse, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_num_split, &v_total_work_items, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_v32_paged_bf16_c80238a3c3::Global()->mod.GetKernel("kernel_mla_decode_v32_paged_bf16"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 231424, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 231424u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 231424u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_v32_paged_bf16_c80238a3c3_1fef86a30b84 + +namespace variant_mla_decode_v32_reduce_1db619cce9_4512a48ed07e { + +void Run(TensorView arg_partial_O, TensorView arg_partial_lse, TensorView arg_O, TensorView arg_sinks, int64_t arg_num_q_heads, int64_t arg_num_split, double arg_bmm2_scale, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_partial_O.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); + CheckContiguous(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) + << "scalar 'num_q_heads' value " << arg_num_q_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) + << "scalar 'num_split' value " << arg_num_split + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_partial_lse, arg_partial_O, "partial_lse", "partial_O"); + CheckSameCudaDevice(arg_O, arg_partial_O, "O", "partial_O"); + CheckSameCudaDevice(arg_sinks, arg_partial_O, "sinks", "partial_O"); + CheckCurrentCudaDevice(arg_partial_O, "partial_O"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + void* p_partial_O = arg_partial_O.data_ptr(); + void* p_partial_lse = arg_partial_lse.data_ptr(); + void* p_O = arg_O.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + int32_t v_num_q_heads = (int32_t)arg_num_q_heads; + int32_t v_num_split = (int32_t)arg_num_split; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_partial_O, &p_partial_lse, &p_O, &p_sinks, &v_num_q_heads, &v_num_split, &v_bmm2_scale, &v_enable_sink}; + + static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(mla_decode_v32_reduce_1db619cce9, "kernel_mla_decode_v32_reduce"); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(128u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[1]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 256u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[1]; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 256u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_v32_reduce_1db619cce9_4512a48ed07e + +int SelectVariant(float softmax_scale_log2, float partial_bmm2_scale, float bmm2_scale, int32_t num_split, int32_t total_work_items, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t num_rows) { + if (((((num_split == 8) && (total_work_items == (num_rows * num_split))) && (num_rows > 0)) && (grid_x > 0))) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_q_rows, tvm::ffi::TensorView arg_kv_half_pages, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sinks, tvm::ffi::TensorView arg_ws_partial_output, tvm::ffi::TensorView arg_ws_partial_lse, float softmax_scale_log2, float partial_bmm2_scale, float bmm2_scale, int32_t num_split, int32_t total_work_items, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t num_rows, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_q_rows, "q_rows"); + CheckCudaTensor(arg_kv_half_pages, "kv_half_pages"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(softmax_scale_log2, partial_bmm2_scale, bmm2_scale, num_split, total_work_items, max_pages_per_seq, enable_sink, grid_x, num_rows); + switch (v) { + case 0: { + CheckCudaTensor(arg_ws_partial_output, "partial_output"); + CheckDtype(arg_ws_partial_output, "partial_output", 4, 16, 1); + TVM_FFI_CHECK(arg_ws_partial_output.numel() == (int64_t)((num_rows) * (128) * (num_split) * (512)), ValueError) + << "workspace 'partial_output' numel " << arg_ws_partial_output.numel() + << " != expected " << (int64_t)((num_rows) * (128) * (num_split) * (512)); + CheckCudaTensor(arg_ws_partial_lse, "partial_lse"); + CheckDtype(arg_ws_partial_lse, "partial_lse", 2, 32, 1); + TVM_FFI_CHECK(arg_ws_partial_lse.numel() == (int64_t)((num_rows) * (128) * (num_split)), ValueError) + << "workspace 'partial_lse' numel " << arg_ws_partial_lse.numel() + << " != expected " << (int64_t)((num_rows) * (128) * (num_split)); + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = 1; + mla_host_shim::variant_mla_decode_v32_paged_bf16_c80238a3c3_1fef86a30b84::Run(arg_q_rows, arg_kv_half_pages, arg_kv_half_pages, arg_kv_half_pages, arg_ws_partial_output, arg_ws_partial_lse, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, partial_bmm2_scale, num_split, total_work_items, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + { + int64_t gx = num_rows; + int64_t gy = 128; + int64_t gz = 1; + mla_host_shim::variant_mla_decode_v32_reduce_1db619cce9_4512a48ed07e::Run(arg_ws_partial_output, arg_ws_partial_lse, arg_output, arg_sinks, 128, num_split, bmm2_scale, enable_sink, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_bf16_native_split8_pdl shape (softmax_scale_log2=" << softmax_scale_log2 << ", partial_bmm2_scale=" << partial_bmm2_scale << ", bmm2_scale=" << bmm2_scale << ", num_split=" << num_split << ", total_work_items=" << total_work_items << ", max_pages_per_seq=" << max_pages_per_seq << ", enable_sink=" << enable_sink << ", grid_x=" << grid_x << ", num_rows=" << num_rows << ")"; + } +} + +int64_t RouteOf(float softmax_scale_log2, float partial_bmm2_scale, float bmm2_scale, int32_t num_split, int32_t total_work_items, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t num_rows) { return SelectVariant(softmax_scale_log2, partial_bmm2_scale, bmm2_scale, num_split, total_work_items, max_pages_per_seq, enable_sink, grid_x, num_rows); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_tail.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_tail.cpp new file mode 100644 index 00000000000..ea53af829b9 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_tail.cpp @@ -0,0 +1,409 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(full_abi_tail_bf16_83451b7317); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +namespace variant_full_abi_tail_bf16_83451b7317_d1d6ab272ce6 { + +void Run(TensorView arg_Q, TensorView arg_KV, TensorView arg_fp8_lut, TensorView arg_page_table, TensorView arg_sparse_indices, TensorView arg_row_batches, TensorView arg_row_seq_lens, TensorView arg_O, TensorView arg_LSE, TensorView arg_sinks, int64_t arg_num_heads, int64_t arg_qk_dim, int64_t arg_value_dim, int64_t arg_kv_stride, int64_t arg_page_size, int64_t arg_page_table_width, int64_t arg_sparse_width, int64_t arg_use_sparse, double arg_softmax_scale, double arg_bmm2_scale, int64_t arg_enable_sink, int64_t arg_write_lse, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_Q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_Q, "Q"); + CheckDtype(arg_Q, "Q", 4, 16, 1); + CheckContiguous(arg_Q, "Q"); + CheckCudaTensor(arg_KV, "KV"); + CheckDtype(arg_KV, "KV", 4, 16, 1); + CheckContiguous(arg_KV, "KV"); + CheckCudaTensor(arg_fp8_lut, "fp8_lut"); + CheckDtype(arg_fp8_lut, "fp8_lut", 2, 32, 1); + CheckContiguous(arg_fp8_lut, "fp8_lut"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sparse_indices, "sparse_indices"); + CheckDtype(arg_sparse_indices, "sparse_indices", 0, 32, 1); + CheckContiguous(arg_sparse_indices, "sparse_indices"); + CheckCudaTensor(arg_row_batches, "row_batches"); + CheckDtype(arg_row_batches, "row_batches", 0, 32, 1); + CheckContiguous(arg_row_batches, "row_batches"); + CheckCudaTensor(arg_row_seq_lens, "row_seq_lens"); + CheckDtype(arg_row_seq_lens, "row_seq_lens", 0, 32, 1); + CheckContiguous(arg_row_seq_lens, "row_seq_lens"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_LSE, "LSE"); + CheckDtype(arg_LSE, "LSE", 2, 32, 1); + CheckContiguous(arg_LSE, "LSE"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_qk_dim >= -2147483648LL && arg_qk_dim <= 2147483647LL, ValueError) + << "scalar 'qk_dim' value " << arg_qk_dim + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_dim >= -2147483648LL && arg_value_dim <= 2147483647LL, ValueError) + << "scalar 'value_dim' value " << arg_value_dim + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_kv_stride >= -2147483648LL && arg_kv_stride <= 2147483647LL, ValueError) + << "scalar 'kv_stride' value " << arg_kv_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_size >= -2147483648LL && arg_page_size <= 2147483647LL, ValueError) + << "scalar 'page_size' value " << arg_page_size + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_table_width >= -2147483648LL && arg_page_table_width <= 2147483647LL, ValueError) + << "scalar 'page_table_width' value " << arg_page_table_width + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_width >= -2147483648LL && arg_sparse_width <= 2147483647LL, ValueError) + << "scalar 'sparse_width' value " << arg_sparse_width + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_sparse >= -2147483648LL && arg_use_sparse <= 2147483647LL, ValueError) + << "scalar 'use_sparse' value " << arg_use_sparse + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_write_lse >= -2147483648LL && arg_write_lse <= 2147483647LL, ValueError) + << "scalar 'write_lse' value " << arg_write_lse + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_KV, arg_Q, "KV", "Q"); + CheckSameCudaDevice(arg_fp8_lut, arg_Q, "fp8_lut", "Q"); + CheckSameCudaDevice(arg_page_table, arg_Q, "page_table", "Q"); + CheckSameCudaDevice(arg_sparse_indices, arg_Q, "sparse_indices", "Q"); + CheckSameCudaDevice(arg_row_batches, arg_Q, "row_batches", "Q"); + CheckSameCudaDevice(arg_row_seq_lens, arg_Q, "row_seq_lens", "Q"); + CheckSameCudaDevice(arg_O, arg_Q, "O", "Q"); + CheckSameCudaDevice(arg_LSE, arg_Q, "LSE", "Q"); + CheckSameCudaDevice(arg_sinks, arg_Q, "sinks", "Q"); + CheckCurrentCudaDevice(arg_Q, "Q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + void* p_Q = arg_Q.data_ptr(); + void* p_KV = arg_KV.data_ptr(); + void* p_fp8_lut = arg_fp8_lut.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sparse_indices = arg_sparse_indices.data_ptr(); + void* p_row_batches = arg_row_batches.data_ptr(); + void* p_row_seq_lens = arg_row_seq_lens.data_ptr(); + void* p_O = arg_O.data_ptr(); + void* p_LSE = arg_LSE.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_qk_dim = (int32_t)arg_qk_dim; + int32_t v_value_dim = (int32_t)arg_value_dim; + int32_t v_kv_stride = (int32_t)arg_kv_stride; + int32_t v_page_size = (int32_t)arg_page_size; + int32_t v_page_table_width = (int32_t)arg_page_table_width; + int32_t v_sparse_width = (int32_t)arg_sparse_width; + int32_t v_use_sparse = (int32_t)arg_use_sparse; + float v_softmax_scale = (float)arg_softmax_scale; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + int32_t v_write_lse = (int32_t)arg_write_lse; + void* kargs[] = {&p_Q, &p_KV, &p_fp8_lut, &p_page_table, &p_sparse_indices, &p_row_batches, &p_row_seq_lens, &p_O, &p_LSE, &p_sinks, &v_num_heads, &v_qk_dim, &v_value_dim, &v_kv_stride, &v_page_size, &v_page_table_width, &v_sparse_width, &v_use_sparse, &v_softmax_scale, &v_bmm2_scale, &v_enable_sink, &v_write_lse}; + + static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(full_abi_tail_bf16_83451b7317, "kernel_full_abi_tail_bf16"); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(32u, 1u, 1u); + + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.Launch(kargs, grid, block, stream, 16384u)); +} +} // namespace variant_full_abi_tail_bf16_83451b7317_d1d6ab272ce6 + +int SelectVariant(int32_t num_heads, int32_t qk_dim, int32_t value_dim, int32_t kv_stride, int32_t page_size, int32_t page_table_width, int32_t sparse_width, int32_t use_sparse, float softmax_scale, float bmm2_scale, int32_t enable_sink, int32_t write_lse, int32_t num_rows) { + if ((num_rows > 0)) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_query, tvm::ffi::TensorView arg_kv, tvm::ffi::TensorView arg_fp8_lut, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sparse_indices, tvm::ffi::TensorView arg_row_batches, tvm::ffi::TensorView arg_row_seq_lens, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_lse, tvm::ffi::TensorView arg_sinks, int32_t num_heads, int32_t qk_dim, int32_t value_dim, int32_t kv_stride, int32_t page_size, int32_t page_table_width, int32_t sparse_width, int32_t use_sparse, float softmax_scale, float bmm2_scale, int32_t enable_sink, int32_t write_lse, int32_t num_rows, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_query, "query"); + CheckCudaTensor(arg_kv, "kv"); + CheckCudaTensor(arg_fp8_lut, "fp8_lut"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sparse_indices, "sparse_indices"); + CheckCudaTensor(arg_row_batches, "row_batches"); + CheckCudaTensor(arg_row_seq_lens, "row_seq_lens"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_lse, "lse"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(num_heads, qk_dim, value_dim, kv_stride, page_size, page_table_width, sparse_width, use_sparse, softmax_scale, bmm2_scale, enable_sink, write_lse, num_rows); + switch (v) { + case 0: { + { + int64_t gx = num_rows; + int64_t gy = num_heads; + int64_t gz = 1; + mla_host_shim::variant_full_abi_tail_bf16_83451b7317_d1d6ab272ce6::Run(arg_query, arg_kv, arg_fp8_lut, arg_page_table, arg_sparse_indices, arg_row_batches, arg_row_seq_lens, arg_output, arg_lse, arg_sinks, num_heads, qk_dim, value_dim, kv_stride, page_size, page_table_width, sparse_width, use_sparse, softmax_scale, bmm2_scale, enable_sink, write_lse, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_bf16_tail shape (num_heads=" << num_heads << ", qk_dim=" << qk_dim << ", value_dim=" << value_dim << ", kv_stride=" << kv_stride << ", page_size=" << page_size << ", page_table_width=" << page_table_width << ", sparse_width=" << sparse_width << ", use_sparse=" << use_sparse << ", softmax_scale=" << softmax_scale << ", bmm2_scale=" << bmm2_scale << ", enable_sink=" << enable_sink << ", write_lse=" << write_lse << ", num_rows=" << num_rows << ")"; + } +} + +int64_t RouteOf(int32_t num_heads, int32_t qk_dim, int32_t value_dim, int32_t kv_stride, int32_t page_size, int32_t page_table_width, int32_t sparse_width, int32_t use_sparse, float softmax_scale, float bmm2_scale, int32_t enable_sink, int32_t write_lse, int32_t num_rows) { return SelectVariant(num_heads, qk_dim, value_dim, kv_stride, page_size, page_table_width, sparse_width, use_sparse, softmax_scale, bmm2_scale, enable_sink, write_lse, num_rows); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_unsplit.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_unsplit.cpp new file mode 100644 index 00000000000..e2a3bc16241 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_unsplit.cpp @@ -0,0 +1,768 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(mla_decode_value_split_paged_bf16_5c36fe8228); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +struct TmaDeviceArena { + static constexpr size_t kSlotsPerChunk = 256; + static constexpr size_t kMaxSlots = 4096; + std::vector chunks; + size_t used = 0; +}; + +// Immutable, process-lifetime device tensor-map slots for the pointer ABI. +// A slot is never rewritten: different descriptor bytes always get a new +// address, so concurrent streams cannot observe a partially updated map. The +// chunked arena caps storage at 512 KiB per CUDA context in this host module. +static inline void* TmaDeviceSlot( + const CUtensorMap& tm, + int device_id, + cudaStream_t stream) { + static std::mutex mu; + static auto* slots = new std::unordered_map(); + static auto* arenas = new std::unordered_map(); + + // Device allocations are context-owned. Resolve and validate the active + // context before cache lookup so a warm entry can never bypass the same + // checks as a cold entry or leak a pointer across contexts on one device. + CUcontext current_context = nullptr; + CUresult result = cuCtxGetCurrent(¤t_context); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_context != nullptr, RuntimeError) + << "pointer TMA ABI requires an active CUDA context: CUresult=" + << static_cast(result); + CUdevice current_device = -1; + result = cuCtxGetDevice(¤t_device); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_device == device_id, RuntimeError) + << "TMA descriptor device mismatch: current=" << current_device + << ", tensor=" << device_id; + + std::string key = + std::to_string(reinterpret_cast(current_context)); + key.push_back(':'); + key.append(reinterpret_cast(&tm), sizeof(CUtensorMap)); + std::lock_guard lock(mu); + + // CUcontext is an opaque, recyclable handle. Before trusting a warm arena, + // prove that its first allocation still belongs to this live context. A + // destroyed context makes the pointer query fail; a subsequently reused + // handle must therefore discard the stale pointer and descriptor entries. + auto arena_it = arenas->find(current_context); + if (arena_it != arenas->end() && !arena_it->second.chunks.empty()) { + CUcontext allocation_context = nullptr; + result = cuPointerGetAttribute( + &allocation_context, + CU_POINTER_ATTRIBUTE_CONTEXT, + arena_it->second.chunks.front()); + if (result != CUDA_SUCCESS || allocation_context != current_context) { + const std::string prefix = + std::to_string(reinterpret_cast(current_context)) + ":"; + for (auto slot_it = slots->begin(); slot_it != slots->end();) { + if (slot_it->first.compare(0, prefix.size(), prefix) == 0) { + slot_it = slots->erase(slot_it); + } else { + ++slot_it; + } + } + arenas->erase(arena_it); + } + } + + auto it = slots->find(key); + if (it != slots->end()) return it->second; + + CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; + result = cuStreamIsCapturing( + reinterpret_cast(stream), &capture_status); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuStreamIsCapturing for TMA descriptor slot failed: CUresult=" + << static_cast(result); + TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) + << "pointer TMA ABI cannot create a new device descriptor slot inside " + "CUDA Graph capture; prewarm this exact tensor/layout binding or " + "compile with tma_abi='grid_constant'"; + + TmaDeviceArena& arena = (*arenas)[current_context]; + TVM_FFI_CHECK(arena.used < TmaDeviceArena::kMaxSlots, RuntimeError) + << "pointer TMA ABI exhausted its immutable descriptor arena in CUDA " + "context " << current_context << " on device " << device_id + << " (capacity=" << TmaDeviceArena::kMaxSlots + << "); reuse tensor/layout bindings or compile with tma_abi='grid_constant'"; + if (arena.used % TmaDeviceArena::kSlotsPerChunk == 0) { + CUdeviceptr chunk = 0; + result = cuMemAlloc( + &chunk, + TmaDeviceArena::kSlotsPerChunk * sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemAlloc for TMA descriptor arena failed: CUresult=" + << static_cast(result); + arena.chunks.push_back(chunk); + } + size_t chunk_index = arena.used / TmaDeviceArena::kSlotsPerChunk; + size_t slot_index = arena.used % TmaDeviceArena::kSlotsPerChunk; + CUdeviceptr dev = arena.chunks[chunk_index] + + slot_index * sizeof(CUtensorMap); + result = cuMemcpyHtoD(dev, &tm, sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemcpyHtoD for TMA descriptor slot failed: CUresult=" + << static_cast(result); + ++arena.used; + void* pointer = reinterpret_cast(static_cast(dev)); + (*slots)[key] = pointer; + return pointer; +} + +namespace variant_mla_decode_value_split_paged_bf16_5c36fe8228_a5bebae701d8 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_nope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_nope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_nope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_nope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_nope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_nope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_nope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_nope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_rope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_rope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_rope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_rope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_rope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_rope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_rope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_rope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k_nope, TensorView arg_tmap_k_rope, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k_nope, "tmap_k_nope"); + CheckDtype(arg_tmap_k_nope, "tmap_k_nope", 4, 16, 1); + CheckContiguous(arg_tmap_k_nope, "tmap_k_nope"); + CheckCudaTensor(arg_tmap_k_rope, "tmap_k_rope"); + CheckDtype(arg_tmap_k_rope, "tmap_k_rope", 4, 16, 1); + CheckContiguous(arg_tmap_k_rope, "tmap_k_rope"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k_nope, arg_tmap_q, "tmap_k_nope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_k_rope, arg_tmap_q, "tmap_k_rope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + void* p_tmap_q = TmaDeviceSlot(h_tmap_q, arg_tmap_q.device().device_id, stream); + CUtensorMap h_tmap_k_nope = EncodeTma_tmap_k_nope(arg_tmap_k_nope); + void* p_tmap_k_nope = TmaDeviceSlot(h_tmap_k_nope, arg_tmap_k_nope.device().device_id, stream); + CUtensorMap h_tmap_k_rope = EncodeTma_tmap_k_rope(arg_tmap_k_rope); + void* p_tmap_k_rope = TmaDeviceSlot(h_tmap_k_rope, arg_tmap_k_rope.device().device_id, stream); + CUtensorMap h_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_tmap_v = TmaDeviceSlot(h_tmap_v, arg_tmap_v.device().device_id, stream); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k_nope, &p_tmap_k_rope, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_value_split_paged_bf16_5c36fe8228::Global()->mod.GetKernel("kernel_mla_decode_value_split_paged_bf16"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 231936, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[2]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 231936u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[2]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 231936u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_value_split_paged_bf16_5c36fe8228_a5bebae701d8 + +int SelectVariant(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z) { + if ((((total_work_items > 0) && (grid_x > 0)) && (grid_z > 0))) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_q_rows, tvm::ffi::TensorView arg_kv_pages, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sinks, float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_q_rows, "q_rows"); + CheckCudaTensor(arg_kv_pages, "kv_pages"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z); + switch (v) { + case 0: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_value_split_paged_bf16_5c36fe8228_a5bebae701d8::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_bf16_unsplit shape (softmax_scale_log2=" << softmax_scale_log2 << ", bmm2_scale=" << bmm2_scale << ", total_work_items=" << total_work_items << ", value_split_count=" << value_split_count << ", max_pages_per_seq=" << max_pages_per_seq << ", enable_sink=" << enable_sink << ", grid_x=" << grid_x << ", grid_z=" << grid_z << ")"; + } +} + +int64_t RouteOf(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z) { return SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_vhalf.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_vhalf.cpp new file mode 100644 index 00000000000..b3b51c6aea6 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_vhalf.cpp @@ -0,0 +1,768 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(mla_decode_value_split_paged_bf16_fa7b59f3a7); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +struct TmaDeviceArena { + static constexpr size_t kSlotsPerChunk = 256; + static constexpr size_t kMaxSlots = 4096; + std::vector chunks; + size_t used = 0; +}; + +// Immutable, process-lifetime device tensor-map slots for the pointer ABI. +// A slot is never rewritten: different descriptor bytes always get a new +// address, so concurrent streams cannot observe a partially updated map. The +// chunked arena caps storage at 512 KiB per CUDA context in this host module. +static inline void* TmaDeviceSlot( + const CUtensorMap& tm, + int device_id, + cudaStream_t stream) { + static std::mutex mu; + static auto* slots = new std::unordered_map(); + static auto* arenas = new std::unordered_map(); + + // Device allocations are context-owned. Resolve and validate the active + // context before cache lookup so a warm entry can never bypass the same + // checks as a cold entry or leak a pointer across contexts on one device. + CUcontext current_context = nullptr; + CUresult result = cuCtxGetCurrent(¤t_context); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_context != nullptr, RuntimeError) + << "pointer TMA ABI requires an active CUDA context: CUresult=" + << static_cast(result); + CUdevice current_device = -1; + result = cuCtxGetDevice(¤t_device); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_device == device_id, RuntimeError) + << "TMA descriptor device mismatch: current=" << current_device + << ", tensor=" << device_id; + + std::string key = + std::to_string(reinterpret_cast(current_context)); + key.push_back(':'); + key.append(reinterpret_cast(&tm), sizeof(CUtensorMap)); + std::lock_guard lock(mu); + + // CUcontext is an opaque, recyclable handle. Before trusting a warm arena, + // prove that its first allocation still belongs to this live context. A + // destroyed context makes the pointer query fail; a subsequently reused + // handle must therefore discard the stale pointer and descriptor entries. + auto arena_it = arenas->find(current_context); + if (arena_it != arenas->end() && !arena_it->second.chunks.empty()) { + CUcontext allocation_context = nullptr; + result = cuPointerGetAttribute( + &allocation_context, + CU_POINTER_ATTRIBUTE_CONTEXT, + arena_it->second.chunks.front()); + if (result != CUDA_SUCCESS || allocation_context != current_context) { + const std::string prefix = + std::to_string(reinterpret_cast(current_context)) + ":"; + for (auto slot_it = slots->begin(); slot_it != slots->end();) { + if (slot_it->first.compare(0, prefix.size(), prefix) == 0) { + slot_it = slots->erase(slot_it); + } else { + ++slot_it; + } + } + arenas->erase(arena_it); + } + } + + auto it = slots->find(key); + if (it != slots->end()) return it->second; + + CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; + result = cuStreamIsCapturing( + reinterpret_cast(stream), &capture_status); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuStreamIsCapturing for TMA descriptor slot failed: CUresult=" + << static_cast(result); + TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) + << "pointer TMA ABI cannot create a new device descriptor slot inside " + "CUDA Graph capture; prewarm this exact tensor/layout binding or " + "compile with tma_abi='grid_constant'"; + + TmaDeviceArena& arena = (*arenas)[current_context]; + TVM_FFI_CHECK(arena.used < TmaDeviceArena::kMaxSlots, RuntimeError) + << "pointer TMA ABI exhausted its immutable descriptor arena in CUDA " + "context " << current_context << " on device " << device_id + << " (capacity=" << TmaDeviceArena::kMaxSlots + << "); reuse tensor/layout bindings or compile with tma_abi='grid_constant'"; + if (arena.used % TmaDeviceArena::kSlotsPerChunk == 0) { + CUdeviceptr chunk = 0; + result = cuMemAlloc( + &chunk, + TmaDeviceArena::kSlotsPerChunk * sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemAlloc for TMA descriptor arena failed: CUresult=" + << static_cast(result); + arena.chunks.push_back(chunk); + } + size_t chunk_index = arena.used / TmaDeviceArena::kSlotsPerChunk; + size_t slot_index = arena.used % TmaDeviceArena::kSlotsPerChunk; + CUdeviceptr dev = arena.chunks[chunk_index] + + slot_index * sizeof(CUtensorMap); + result = cuMemcpyHtoD(dev, &tm, sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemcpyHtoD for TMA descriptor slot failed: CUresult=" + << static_cast(result); + ++arena.used; + void* pointer = reinterpret_cast(static_cast(dev)); + (*slots)[key] = pointer; + return pointer; +} + +namespace variant_mla_decode_value_split_paged_bf16_fa7b59f3a7_ed851bc7fbe4 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_nope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_nope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_nope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_nope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_nope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_nope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_nope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_nope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_rope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_rope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_rope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_rope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_rope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_rope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_rope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_rope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k_nope, TensorView arg_tmap_k_rope, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k_nope, "tmap_k_nope"); + CheckDtype(arg_tmap_k_nope, "tmap_k_nope", 4, 16, 1); + CheckContiguous(arg_tmap_k_nope, "tmap_k_nope"); + CheckCudaTensor(arg_tmap_k_rope, "tmap_k_rope"); + CheckDtype(arg_tmap_k_rope, "tmap_k_rope", 4, 16, 1); + CheckContiguous(arg_tmap_k_rope, "tmap_k_rope"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k_nope, arg_tmap_q, "tmap_k_nope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_k_rope, arg_tmap_q, "tmap_k_rope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + void* p_tmap_q = TmaDeviceSlot(h_tmap_q, arg_tmap_q.device().device_id, stream); + CUtensorMap h_tmap_k_nope = EncodeTma_tmap_k_nope(arg_tmap_k_nope); + void* p_tmap_k_nope = TmaDeviceSlot(h_tmap_k_nope, arg_tmap_k_nope.device().device_id, stream); + CUtensorMap h_tmap_k_rope = EncodeTma_tmap_k_rope(arg_tmap_k_rope); + void* p_tmap_k_rope = TmaDeviceSlot(h_tmap_k_rope, arg_tmap_k_rope.device().device_id, stream); + CUtensorMap h_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_tmap_v = TmaDeviceSlot(h_tmap_v, arg_tmap_v.device().device_id, stream); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k_nope, &p_tmap_k_rope, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_value_split_paged_bf16_fa7b59f3a7::Global()->mod.GetKernel("kernel_mla_decode_value_split_paged_bf16"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 231424, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[2]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 231424u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[2]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 231424u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_value_split_paged_bf16_fa7b59f3a7_ed851bc7fbe4 + +int SelectVariant(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z) { + if ((((total_work_items > 0) && (grid_x > 0)) && (grid_z > 0))) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_q_rows, tvm::ffi::TensorView arg_kv_pages, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sinks, float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_q_rows, "q_rows"); + CheckCudaTensor(arg_kv_pages, "kv_pages"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z); + switch (v) { + case 0: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_value_split_paged_bf16_fa7b59f3a7_ed851bc7fbe4::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_bf16_vhalf shape (softmax_scale_log2=" << softmax_scale_log2 << ", bmm2_scale=" << bmm2_scale << ", total_work_items=" << total_work_items << ", value_split_count=" << value_split_count << ", max_pages_per_seq=" << max_pages_per_seq << ", enable_sink=" << enable_sink << ", grid_x=" << grid_x << ", grid_z=" << grid_z << ")"; + } +} + +int64_t RouteOf(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z) { return SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_vquarter.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_vquarter.cpp new file mode 100644 index 00000000000..bf2e5f73ed2 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_bf16_vquarter.cpp @@ -0,0 +1,768 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(mla_decode_value_split_paged_bf16_ee41d3ceea); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +struct TmaDeviceArena { + static constexpr size_t kSlotsPerChunk = 256; + static constexpr size_t kMaxSlots = 4096; + std::vector chunks; + size_t used = 0; +}; + +// Immutable, process-lifetime device tensor-map slots for the pointer ABI. +// A slot is never rewritten: different descriptor bytes always get a new +// address, so concurrent streams cannot observe a partially updated map. The +// chunked arena caps storage at 512 KiB per CUDA context in this host module. +static inline void* TmaDeviceSlot( + const CUtensorMap& tm, + int device_id, + cudaStream_t stream) { + static std::mutex mu; + static auto* slots = new std::unordered_map(); + static auto* arenas = new std::unordered_map(); + + // Device allocations are context-owned. Resolve and validate the active + // context before cache lookup so a warm entry can never bypass the same + // checks as a cold entry or leak a pointer across contexts on one device. + CUcontext current_context = nullptr; + CUresult result = cuCtxGetCurrent(¤t_context); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_context != nullptr, RuntimeError) + << "pointer TMA ABI requires an active CUDA context: CUresult=" + << static_cast(result); + CUdevice current_device = -1; + result = cuCtxGetDevice(¤t_device); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_device == device_id, RuntimeError) + << "TMA descriptor device mismatch: current=" << current_device + << ", tensor=" << device_id; + + std::string key = + std::to_string(reinterpret_cast(current_context)); + key.push_back(':'); + key.append(reinterpret_cast(&tm), sizeof(CUtensorMap)); + std::lock_guard lock(mu); + + // CUcontext is an opaque, recyclable handle. Before trusting a warm arena, + // prove that its first allocation still belongs to this live context. A + // destroyed context makes the pointer query fail; a subsequently reused + // handle must therefore discard the stale pointer and descriptor entries. + auto arena_it = arenas->find(current_context); + if (arena_it != arenas->end() && !arena_it->second.chunks.empty()) { + CUcontext allocation_context = nullptr; + result = cuPointerGetAttribute( + &allocation_context, + CU_POINTER_ATTRIBUTE_CONTEXT, + arena_it->second.chunks.front()); + if (result != CUDA_SUCCESS || allocation_context != current_context) { + const std::string prefix = + std::to_string(reinterpret_cast(current_context)) + ":"; + for (auto slot_it = slots->begin(); slot_it != slots->end();) { + if (slot_it->first.compare(0, prefix.size(), prefix) == 0) { + slot_it = slots->erase(slot_it); + } else { + ++slot_it; + } + } + arenas->erase(arena_it); + } + } + + auto it = slots->find(key); + if (it != slots->end()) return it->second; + + CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; + result = cuStreamIsCapturing( + reinterpret_cast(stream), &capture_status); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuStreamIsCapturing for TMA descriptor slot failed: CUresult=" + << static_cast(result); + TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) + << "pointer TMA ABI cannot create a new device descriptor slot inside " + "CUDA Graph capture; prewarm this exact tensor/layout binding or " + "compile with tma_abi='grid_constant'"; + + TmaDeviceArena& arena = (*arenas)[current_context]; + TVM_FFI_CHECK(arena.used < TmaDeviceArena::kMaxSlots, RuntimeError) + << "pointer TMA ABI exhausted its immutable descriptor arena in CUDA " + "context " << current_context << " on device " << device_id + << " (capacity=" << TmaDeviceArena::kMaxSlots + << "); reuse tensor/layout bindings or compile with tma_abi='grid_constant'"; + if (arena.used % TmaDeviceArena::kSlotsPerChunk == 0) { + CUdeviceptr chunk = 0; + result = cuMemAlloc( + &chunk, + TmaDeviceArena::kSlotsPerChunk * sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemAlloc for TMA descriptor arena failed: CUresult=" + << static_cast(result); + arena.chunks.push_back(chunk); + } + size_t chunk_index = arena.used / TmaDeviceArena::kSlotsPerChunk; + size_t slot_index = arena.used % TmaDeviceArena::kSlotsPerChunk; + CUdeviceptr dev = arena.chunks[chunk_index] + + slot_index * sizeof(CUtensorMap); + result = cuMemcpyHtoD(dev, &tm, sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemcpyHtoD for TMA descriptor slot failed: CUresult=" + << static_cast(result); + ++arena.used; + void* pointer = reinterpret_cast(static_cast(dev)); + (*slots)[key] = pointer; + return pointer; +} + +namespace variant_mla_decode_value_split_paged_bf16_ee41d3ceea_3d9d1dc71fbc { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_nope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_nope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_nope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_nope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_nope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_nope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_nope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_nope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_rope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_rope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_rope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_rope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_rope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_rope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_rope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_rope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k_nope, TensorView arg_tmap_k_rope, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k_nope, "tmap_k_nope"); + CheckDtype(arg_tmap_k_nope, "tmap_k_nope", 4, 16, 1); + CheckContiguous(arg_tmap_k_nope, "tmap_k_nope"); + CheckCudaTensor(arg_tmap_k_rope, "tmap_k_rope"); + CheckDtype(arg_tmap_k_rope, "tmap_k_rope", 4, 16, 1); + CheckContiguous(arg_tmap_k_rope, "tmap_k_rope"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k_nope, arg_tmap_q, "tmap_k_nope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_k_rope, arg_tmap_q, "tmap_k_rope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + void* p_tmap_q = TmaDeviceSlot(h_tmap_q, arg_tmap_q.device().device_id, stream); + CUtensorMap h_tmap_k_nope = EncodeTma_tmap_k_nope(arg_tmap_k_nope); + void* p_tmap_k_nope = TmaDeviceSlot(h_tmap_k_nope, arg_tmap_k_nope.device().device_id, stream); + CUtensorMap h_tmap_k_rope = EncodeTma_tmap_k_rope(arg_tmap_k_rope); + void* p_tmap_k_rope = TmaDeviceSlot(h_tmap_k_rope, arg_tmap_k_rope.device().device_id, stream); + CUtensorMap h_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_tmap_v = TmaDeviceSlot(h_tmap_v, arg_tmap_v.device().device_id, stream); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k_nope, &p_tmap_k_rope, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_value_split_paged_bf16_ee41d3ceea::Global()->mod.GetKernel("kernel_mla_decode_value_split_paged_bf16"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 231424, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[2]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 231424u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[2]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 231424u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_value_split_paged_bf16_ee41d3ceea_3d9d1dc71fbc + +int SelectVariant(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z) { + if ((((total_work_items > 0) && (grid_x > 0)) && (grid_z > 0))) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_q_rows, tvm::ffi::TensorView arg_kv_pages, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sinks, float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_q_rows, "q_rows"); + CheckCudaTensor(arg_kv_pages, "kv_pages"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z); + switch (v) { + case 0: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_value_split_paged_bf16_ee41d3ceea_3d9d1dc71fbc::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_bf16_vquarter shape (softmax_scale_log2=" << softmax_scale_log2 << ", bmm2_scale=" << bmm2_scale << ", total_work_items=" << total_work_items << ", value_split_count=" << value_split_count << ", max_pages_per_seq=" << max_pages_per_seq << ", enable_sink=" << enable_sink << ", grid_x=" << grid_x << ", grid_z=" << grid_z << ")"; + } +} + +int64_t RouteOf(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z) { return SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_fp8_p32_qk_l2.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_fp8_p32_qk_l2.cpp new file mode 100644 index 00000000000..2a0315e2c32 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_fp8_p32_qk_l2.cpp @@ -0,0 +1,560 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(fp8_p32_exact_producer_58f57b1b1c); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +namespace variant_fp8_p32_exact_producer_58f57b1b1c_6d7596abb89c { + +// 2D TMA descriptor for buffer 'Q_map' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_Q_map(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'Q_map' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'Q_map' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'Q_map' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "Q_map"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'Q_map' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'Q_map' resolved a non-positive global dim"; + TVM_FFI_CHECK(16u <= global_dim[1], ValueError) + << "TMA box (128, 16) exceeds resolved global dims for 'Q_map'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'Q_map' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'Q_map' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'Q_map' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 16u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'Q_map') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'K_map' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_K_map(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'K_map' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'K_map' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'K_map' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "K_map"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'K_map' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'K_map' resolved a non-positive global dim"; + TVM_FFI_CHECK(32u <= global_dim[1], ValueError) + << "TMA box (128, 32) exceeds resolved global dims for 'K_map'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'K_map' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'K_map' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'K_map' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 32u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'K_map') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'V_map' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_V_map(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'V_map' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'V_map' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'V_map' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "V_map"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'V_map' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'V_map' resolved a non-positive global dim"; + TVM_FFI_CHECK(32u <= global_dim[1], ValueError) + << "TMA box (128, 32) exceeds resolved global dims for 'V_map'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'V_map' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'V_map' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'V_map' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 32u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'V_map') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_Q_map, TensorView arg_K_map, TensorView arg_V_map, TensorView arg_page_table, TensorView arg_seq_lens, TensorView arg_partial_O, TensorView arg_partial_stats, TensorView arg_O, TensorView arg_completion, int64_t arg_q_len, int64_t arg_page_table_stride, int64_t arg_max_num_ctas_q, int64_t arg_max_num_ctas_kv, double arg_bmm1_scale_log2, double arg_bmm2_scale, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_Q_map.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_Q_map, "Q_map"); + CheckDtype(arg_Q_map, "Q_map", 1, 8, 1); + CheckCudaTensor(arg_K_map, "K_map"); + CheckDtype(arg_K_map, "K_map", 1, 8, 1); + CheckCudaTensor(arg_V_map, "V_map"); + CheckDtype(arg_V_map, "V_map", 1, 8, 1); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckDtype(arg_seq_lens, "seq_lens", 0, 32, 1); + CheckContiguous(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); + CheckContiguous(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_stats, "partial_stats"); + CheckDtype(arg_partial_stats, "partial_stats", 2, 32, 1); + CheckContiguous(arg_partial_stats, "partial_stats"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_completion, "completion"); + CheckDtype(arg_completion, "completion", 1, 32, 1); + CheckContiguous(arg_completion, "completion"); + TVM_FFI_CHECK(arg_q_len >= -2147483648LL && arg_q_len <= 2147483647LL, ValueError) + << "scalar 'q_len' value " << arg_q_len + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_table_stride >= -2147483648LL && arg_page_table_stride <= 2147483647LL, ValueError) + << "scalar 'page_table_stride' value " << arg_page_table_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_num_ctas_q >= -2147483648LL && arg_max_num_ctas_q <= 2147483647LL, ValueError) + << "scalar 'max_num_ctas_q' value " << arg_max_num_ctas_q + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_num_ctas_kv >= -2147483648LL && arg_max_num_ctas_kv <= 2147483647LL, ValueError) + << "scalar 'max_num_ctas_kv' value " << arg_max_num_ctas_kv + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_K_map, arg_Q_map, "K_map", "Q_map"); + CheckSameCudaDevice(arg_V_map, arg_Q_map, "V_map", "Q_map"); + CheckSameCudaDevice(arg_page_table, arg_Q_map, "page_table", "Q_map"); + CheckSameCudaDevice(arg_seq_lens, arg_Q_map, "seq_lens", "Q_map"); + CheckSameCudaDevice(arg_partial_O, arg_Q_map, "partial_O", "Q_map"); + CheckSameCudaDevice(arg_partial_stats, arg_Q_map, "partial_stats", "Q_map"); + CheckSameCudaDevice(arg_O, arg_Q_map, "O", "Q_map"); + CheckSameCudaDevice(arg_completion, arg_Q_map, "completion", "Q_map"); + CheckCurrentCudaDevice(arg_Q_map, "Q_map"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + CUtensorMap p_Q_map = EncodeTma_Q_map(arg_Q_map); + CUtensorMap p_K_map = EncodeTma_K_map(arg_K_map); + CUtensorMap p_V_map = EncodeTma_V_map(arg_V_map); + void* p_page_table = arg_page_table.data_ptr(); + void* p_seq_lens = arg_seq_lens.data_ptr(); + void* p_partial_O = arg_partial_O.data_ptr(); + void* p_partial_stats = arg_partial_stats.data_ptr(); + void* p_O = arg_O.data_ptr(); + void* p_completion = arg_completion.data_ptr(); + int32_t v_q_len = (int32_t)arg_q_len; + int32_t v_page_table_stride = (int32_t)arg_page_table_stride; + int32_t v_max_num_ctas_q = (int32_t)arg_max_num_ctas_q; + int32_t v_max_num_ctas_kv = (int32_t)arg_max_num_ctas_kv; + float v_bmm1_scale_log2 = (float)arg_bmm1_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + void* kargs[] = {&p_Q_map, &p_K_map, &p_V_map, &p_page_table, &p_seq_lens, &p_partial_O, &p_partial_stats, &p_O, &p_completion, &v_q_len, &v_page_table_stride, &v_max_num_ctas_q, &v_max_num_ctas_kv, &v_bmm1_scale_log2, &v_bmm2_scale}; + + static auto kernel = EmbedCubinModule_fp8_p32_exact_producer_58f57b1b1c::Global()->mod.GetKernel("kernel_fp8_p32_exact_producer"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_Q_map.device().device_id, 178560, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(512u, 1u, 1u); + + if (use_oversized_smem) { + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[1]; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 178560u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[1]; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 178560u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); + + } else { + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.Launch(kargs, grid, block, stream, 178560u)); + } +} +} // namespace variant_fp8_p32_exact_producer_58f57b1b1c_6d7596abb89c + +int SelectVariant(int32_t batch, int32_t q_len, int32_t page_table_stride, int32_t max_num_ctas_q, int32_t max_num_ctas_kv, float bmm1_scale_log2, float bmm2_scale) { + if (((batch > 0) && (q_len > 0))) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_query, tvm::ffi::TensorView arg_kv, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_completion, tvm::ffi::TensorView arg_ws_partial_output, tvm::ffi::TensorView arg_ws_partial_stats, int32_t batch, int32_t q_len, int32_t page_table_stride, int32_t max_num_ctas_q, int32_t max_num_ctas_kv, float bmm1_scale_log2, float bmm2_scale, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_query, "query"); + CheckCudaTensor(arg_kv, "kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_completion, "completion"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(batch, q_len, page_table_stride, max_num_ctas_q, max_num_ctas_kv, bmm1_scale_log2, bmm2_scale); + switch (v) { + case 0: { + CheckCudaTensor(arg_ws_partial_output, "partial_output"); + CheckDtype(arg_ws_partial_output, "partial_output", 4, 16, 1); + TVM_FFI_CHECK(arg_ws_partial_output.numel() == (int64_t)((((batch * 32) * q_len)) * (2) * (16) * (128)), ValueError) + << "workspace 'partial_output' numel " << arg_ws_partial_output.numel() + << " != expected " << (int64_t)((((batch * 32) * q_len)) * (2) * (16) * (128)); + CheckCudaTensor(arg_ws_partial_stats, "partial_stats"); + CheckDtype(arg_ws_partial_stats, "partial_stats", 2, 32, 1); + TVM_FFI_CHECK(arg_ws_partial_stats.numel() == (int64_t)((((batch * 32) * q_len)) * (2) * (16) * (2)), ValueError) + << "workspace 'partial_stats' numel " << arg_ws_partial_stats.numel() + << " != expected " << (int64_t)((((batch * 32) * q_len)) * (2) * (16) * (2)); + { + int64_t gx = (q_len * 2); + int64_t gy = 32; + int64_t gz = batch; + mla_host_shim::variant_fp8_p32_exact_producer_58f57b1b1c_6d7596abb89c::Run(arg_query, arg_kv, arg_kv, arg_page_table, arg_seq_lens, arg_ws_partial_output, arg_ws_partial_stats, arg_output, arg_completion, q_len, page_table_stride, max_num_ctas_q, max_num_ctas_kv, bmm1_scale_log2, bmm2_scale, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_fp8_p32_qk_l2 shape (batch=" << batch << ", q_len=" << q_len << ", page_table_stride=" << page_table_stride << ", max_num_ctas_q=" << max_num_ctas_q << ", max_num_ctas_kv=" << max_num_ctas_kv << ", bmm1_scale_log2=" << bmm1_scale_log2 << ", bmm2_scale=" << bmm2_scale << ")"; + } +} + +int64_t RouteOf(int32_t batch, int32_t q_len, int32_t page_table_stride, int32_t max_num_ctas_q, int32_t max_num_ctas_kv, float bmm1_scale_log2, float bmm2_scale) { return SelectVariant(batch, q_len, page_table_stride, max_num_ctas_q, max_num_ctas_kv, bmm1_scale_log2, bmm2_scale); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_fp8_page64_pdl.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_fp8_page64_pdl.cpp new file mode 100644 index 00000000000..425a4ba4928 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_fp8_page64_pdl.cpp @@ -0,0 +1,767 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(page64_keeps_gmem_separate_reduce_cea01a87e9); +TVM_FFI_EMBED_CUBIN(trtllm_mla_decode_fp8_06282d4d84); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +struct TmaDeviceArena { + static constexpr size_t kSlotsPerChunk = 256; + static constexpr size_t kMaxSlots = 4096; + std::vector chunks; + size_t used = 0; +}; + +// Immutable, process-lifetime device tensor-map slots for the pointer ABI. +// A slot is never rewritten: different descriptor bytes always get a new +// address, so concurrent streams cannot observe a partially updated map. The +// chunked arena caps storage at 512 KiB per CUDA context in this host module. +static inline void* TmaDeviceSlot( + const CUtensorMap& tm, + int device_id, + cudaStream_t stream) { + static std::mutex mu; + static auto* slots = new std::unordered_map(); + static auto* arenas = new std::unordered_map(); + + // Device allocations are context-owned. Resolve and validate the active + // context before cache lookup so a warm entry can never bypass the same + // checks as a cold entry or leak a pointer across contexts on one device. + CUcontext current_context = nullptr; + CUresult result = cuCtxGetCurrent(¤t_context); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_context != nullptr, RuntimeError) + << "pointer TMA ABI requires an active CUDA context: CUresult=" + << static_cast(result); + CUdevice current_device = -1; + result = cuCtxGetDevice(¤t_device); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_device == device_id, RuntimeError) + << "TMA descriptor device mismatch: current=" << current_device + << ", tensor=" << device_id; + + std::string key = + std::to_string(reinterpret_cast(current_context)); + key.push_back(':'); + key.append(reinterpret_cast(&tm), sizeof(CUtensorMap)); + std::lock_guard lock(mu); + + // CUcontext is an opaque, recyclable handle. Before trusting a warm arena, + // prove that its first allocation still belongs to this live context. A + // destroyed context makes the pointer query fail; a subsequently reused + // handle must therefore discard the stale pointer and descriptor entries. + auto arena_it = arenas->find(current_context); + if (arena_it != arenas->end() && !arena_it->second.chunks.empty()) { + CUcontext allocation_context = nullptr; + result = cuPointerGetAttribute( + &allocation_context, + CU_POINTER_ATTRIBUTE_CONTEXT, + arena_it->second.chunks.front()); + if (result != CUDA_SUCCESS || allocation_context != current_context) { + const std::string prefix = + std::to_string(reinterpret_cast(current_context)) + ":"; + for (auto slot_it = slots->begin(); slot_it != slots->end();) { + if (slot_it->first.compare(0, prefix.size(), prefix) == 0) { + slot_it = slots->erase(slot_it); + } else { + ++slot_it; + } + } + arenas->erase(arena_it); + } + } + + auto it = slots->find(key); + if (it != slots->end()) return it->second; + + CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; + result = cuStreamIsCapturing( + reinterpret_cast(stream), &capture_status); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuStreamIsCapturing for TMA descriptor slot failed: CUresult=" + << static_cast(result); + TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) + << "pointer TMA ABI cannot create a new device descriptor slot inside " + "CUDA Graph capture; prewarm this exact tensor/layout binding or " + "compile with tma_abi='grid_constant'"; + + TmaDeviceArena& arena = (*arenas)[current_context]; + TVM_FFI_CHECK(arena.used < TmaDeviceArena::kMaxSlots, RuntimeError) + << "pointer TMA ABI exhausted its immutable descriptor arena in CUDA " + "context " << current_context << " on device " << device_id + << " (capacity=" << TmaDeviceArena::kMaxSlots + << "); reuse tensor/layout bindings or compile with tma_abi='grid_constant'"; + if (arena.used % TmaDeviceArena::kSlotsPerChunk == 0) { + CUdeviceptr chunk = 0; + result = cuMemAlloc( + &chunk, + TmaDeviceArena::kSlotsPerChunk * sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemAlloc for TMA descriptor arena failed: CUresult=" + << static_cast(result); + arena.chunks.push_back(chunk); + } + size_t chunk_index = arena.used / TmaDeviceArena::kSlotsPerChunk; + size_t slot_index = arena.used % TmaDeviceArena::kSlotsPerChunk; + CUdeviceptr dev = arena.chunks[chunk_index] + + slot_index * sizeof(CUtensorMap); + result = cuMemcpyHtoD(dev, &tm, sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemcpyHtoD for TMA descriptor slot failed: CUresult=" + << static_cast(result); + ++arena.used; + void* pointer = reinterpret_cast(static_cast(dev)); + (*slots)[key] = pointer; + return pointer; +} + +namespace variant_trtllm_mla_decode_fp8_06282d4d84_38925e5a8bfb { + +// 2D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'tmap_q' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[1], ValueError) + << "TMA box (128, 64) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 64u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'tmap_k' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[1], ValueError) + << "TMA box (128, 64) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 64u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'tmap_v' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[1], ValueError) + << "TMA box (128, 64) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 64u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_partial_O, TensorView arg_partial_stats, TensorView arg_O, TensorView arg_LSE, TensorView arg_seq_lens_kv, TensorView arg_work_batch_indices, TensorView arg_page_table, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_num_split, int64_t arg_num_queries, int64_t arg_max_pages_per_seq, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 1, 8, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 1, 8, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); + CheckContiguous(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_stats, "partial_stats"); + CheckDtype(arg_partial_stats, "partial_stats", 2, 32, 1); + CheckContiguous(arg_partial_stats, "partial_stats"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_LSE, "LSE"); + CheckDtype(arg_LSE, "LSE", 2, 32, 1); + CheckContiguous(arg_LSE, "LSE"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_work_batch_indices, "work_batch_indices"); + CheckDtype(arg_work_batch_indices, "work_batch_indices", 0, 32, 1); + CheckContiguous(arg_work_batch_indices, "work_batch_indices"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) + << "scalar 'num_split' value " << arg_num_split + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_queries >= -2147483648LL && arg_num_queries <= 2147483647LL, ValueError) + << "scalar 'num_queries' value " << arg_num_queries + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_partial_O, arg_tmap_q, "partial_O", "tmap_q"); + CheckSameCudaDevice(arg_partial_stats, arg_tmap_q, "partial_stats", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_LSE, arg_tmap_q, "LSE", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_work_batch_indices, arg_tmap_q, "work_batch_indices", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + void* p_tmap_q = TmaDeviceSlot(h_tmap_q, arg_tmap_q.device().device_id, stream); + CUtensorMap h_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + void* p_tmap_k = TmaDeviceSlot(h_tmap_k, arg_tmap_k.device().device_id, stream); + CUtensorMap h_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_tmap_v = TmaDeviceSlot(h_tmap_v, arg_tmap_v.device().device_id, stream); + void* p_partial_O = arg_partial_O.data_ptr(); + void* p_partial_stats = arg_partial_stats.data_ptr(); + void* p_O = arg_O.data_ptr(); + void* p_LSE = arg_LSE.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_work_batch_indices = arg_work_batch_indices.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_num_split = (int32_t)arg_num_split; + int32_t v_num_queries = (int32_t)arg_num_queries; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_partial_O, &p_partial_stats, &p_O, &p_LSE, &p_seq_lens_kv, &p_work_batch_indices, &p_page_table, &v_softmax_scale_log2, &v_bmm2_scale, &v_num_split, &v_num_queries, &v_max_pages_per_seq}; + + static auto kernel = EmbedCubinModule_trtllm_mla_decode_fp8_06282d4d84::Global()->mod.GetKernel("kernel_trtllm_mla_decode_fp8"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 226304, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 226304u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 226304u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_trtllm_mla_decode_fp8_06282d4d84_38925e5a8bfb + +namespace variant_page64_keeps_gmem_separate_reduce_cea01a87e9_dde077f61656 { + +void Run(TensorView arg_partial_O, TensorView arg_partial_stats, TensorView arg_O, TensorView arg_LSE, int64_t arg_num_split, int64_t arg_num_reduce_ctas, int64_t arg_num_queries, double arg_bmm2_scale, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_partial_O.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); + CheckContiguous(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_stats, "partial_stats"); + CheckDtype(arg_partial_stats, "partial_stats", 2, 32, 1); + CheckContiguous(arg_partial_stats, "partial_stats"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_LSE, "LSE"); + CheckDtype(arg_LSE, "LSE", 2, 32, 1); + CheckContiguous(arg_LSE, "LSE"); + TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) + << "scalar 'num_split' value " << arg_num_split + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_reduce_ctas >= -2147483648LL && arg_num_reduce_ctas <= 2147483647LL, ValueError) + << "scalar 'num_reduce_ctas' value " << arg_num_reduce_ctas + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_queries >= -2147483648LL && arg_num_queries <= 2147483647LL, ValueError) + << "scalar 'num_queries' value " << arg_num_queries + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_partial_stats, arg_partial_O, "partial_stats", "partial_O"); + CheckSameCudaDevice(arg_O, arg_partial_O, "O", "partial_O"); + CheckSameCudaDevice(arg_LSE, arg_partial_O, "LSE", "partial_O"); + CheckCurrentCudaDevice(arg_partial_O, "partial_O"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + void* p_partial_O = arg_partial_O.data_ptr(); + void* p_partial_stats = arg_partial_stats.data_ptr(); + void* p_O = arg_O.data_ptr(); + void* p_LSE = arg_LSE.data_ptr(); + int32_t v_num_split = (int32_t)arg_num_split; + int32_t v_num_reduce_ctas = (int32_t)arg_num_reduce_ctas; + int32_t v_num_queries = (int32_t)arg_num_queries; + float v_bmm2_scale = (float)arg_bmm2_scale; + void* kargs[] = {&p_partial_O, &p_partial_stats, &p_O, &p_LSE, &v_num_split, &v_num_reduce_ctas, &v_num_queries, &v_bmm2_scale}; + + static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(page64_keeps_gmem_separate_reduce_cea01a87e9, "kernel_page64_keeps_gmem_separate_reduce"); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(512u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[1]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 0u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[1]; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 0u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_page64_keeps_gmem_separate_reduce_cea01a87e9_dde077f61656 + +int SelectVariant(float softmax_scale_log2, float bmm2_scale, int32_t num_split, int32_t num_queries, int32_t max_pages_per_seq, int32_t num_reduce_ctas) { + if ((num_split > 1)) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_query, tvm::ffi::TensorView arg_kv, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_lse, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_work_batch_indices, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_ws_partial_output, tvm::ffi::TensorView arg_ws_partial_stats, float softmax_scale_log2, float bmm2_scale, int32_t num_split, int32_t num_queries, int32_t max_pages_per_seq, int32_t num_reduce_ctas, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_query, "query"); + CheckCudaTensor(arg_kv, "kv"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_lse, "lse"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_work_batch_indices, "work_batch_indices"); + CheckCudaTensor(arg_page_table, "page_table"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(softmax_scale_log2, bmm2_scale, num_split, num_queries, max_pages_per_seq, num_reduce_ctas); + switch (v) { + case 0: { + CheckCudaTensor(arg_ws_partial_output, "partial_output"); + CheckDtype(arg_ws_partial_output, "partial_output", 4, 16, 1); + TVM_FFI_CHECK(arg_ws_partial_output.numel() == (int64_t)((num_queries) * (num_split) * (128) * (512)), ValueError) + << "workspace 'partial_output' numel " << arg_ws_partial_output.numel() + << " != expected " << (int64_t)((num_queries) * (num_split) * (128) * (512)); + CheckCudaTensor(arg_ws_partial_stats, "partial_stats"); + CheckDtype(arg_ws_partial_stats, "partial_stats", 2, 32, 1); + TVM_FFI_CHECK(arg_ws_partial_stats.numel() == (int64_t)((num_queries) * (num_split) * (128) * (2)), ValueError) + << "workspace 'partial_stats' numel " << arg_ws_partial_stats.numel() + << " != expected " << (int64_t)((num_queries) * (num_split) * (128) * (2)); + { + int64_t gx = (num_split * 2); + int64_t gy = num_queries; + int64_t gz = 1; + mla_host_shim::variant_trtllm_mla_decode_fp8_06282d4d84_38925e5a8bfb::Run(arg_query, arg_kv, arg_kv, arg_ws_partial_output, arg_ws_partial_stats, arg_output, arg_lse, arg_seq_lens, arg_work_batch_indices, arg_page_table, softmax_scale_log2, bmm2_scale, num_split, num_queries, max_pages_per_seq, gx, gy, gz, stream); + } + { + int64_t gx = num_reduce_ctas; + int64_t gy = 2; + int64_t gz = num_queries; + mla_host_shim::variant_page64_keeps_gmem_separate_reduce_cea01a87e9_dde077f61656::Run(arg_ws_partial_output, arg_ws_partial_stats, arg_output, arg_lse, num_split, num_reduce_ctas, num_queries, bmm2_scale, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_fp8_page64_pdl shape (softmax_scale_log2=" << softmax_scale_log2 << ", bmm2_scale=" << bmm2_scale << ", num_split=" << num_split << ", num_queries=" << num_queries << ", max_pages_per_seq=" << max_pages_per_seq << ", num_reduce_ctas=" << num_reduce_ctas << ")"; + } +} + +int64_t RouteOf(float softmax_scale_log2, float bmm2_scale, int32_t num_split, int32_t num_queries, int32_t max_pages_per_seq, int32_t num_reduce_ctas) { return SelectVariant(softmax_scale_log2, bmm2_scale, num_split, num_queries, max_pages_per_seq, num_reduce_ctas); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_fp8_tail.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_fp8_tail.cpp new file mode 100644 index 00000000000..f551a383035 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_100a_148/cake_mla_fp8_tail.cpp @@ -0,0 +1,409 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(full_abi_tail_fp8_26b36022eb); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +namespace variant_full_abi_tail_fp8_26b36022eb_f0c2469a092e { + +void Run(TensorView arg_Q, TensorView arg_KV, TensorView arg_fp8_lut, TensorView arg_page_table, TensorView arg_sparse_indices, TensorView arg_row_batches, TensorView arg_row_seq_lens, TensorView arg_O, TensorView arg_LSE, TensorView arg_sinks, int64_t arg_num_heads, int64_t arg_qk_dim, int64_t arg_value_dim, int64_t arg_kv_stride, int64_t arg_page_size, int64_t arg_page_table_width, int64_t arg_sparse_width, int64_t arg_use_sparse, double arg_softmax_scale, double arg_bmm2_scale, int64_t arg_enable_sink, int64_t arg_write_lse, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_Q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_Q, "Q"); + CheckDtype(arg_Q, "Q", 1, 8, 1); + CheckContiguous(arg_Q, "Q"); + CheckCudaTensor(arg_KV, "KV"); + CheckDtype(arg_KV, "KV", 1, 8, 1); + CheckContiguous(arg_KV, "KV"); + CheckCudaTensor(arg_fp8_lut, "fp8_lut"); + CheckDtype(arg_fp8_lut, "fp8_lut", 2, 32, 1); + CheckContiguous(arg_fp8_lut, "fp8_lut"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sparse_indices, "sparse_indices"); + CheckDtype(arg_sparse_indices, "sparse_indices", 0, 32, 1); + CheckContiguous(arg_sparse_indices, "sparse_indices"); + CheckCudaTensor(arg_row_batches, "row_batches"); + CheckDtype(arg_row_batches, "row_batches", 0, 32, 1); + CheckContiguous(arg_row_batches, "row_batches"); + CheckCudaTensor(arg_row_seq_lens, "row_seq_lens"); + CheckDtype(arg_row_seq_lens, "row_seq_lens", 0, 32, 1); + CheckContiguous(arg_row_seq_lens, "row_seq_lens"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_LSE, "LSE"); + CheckDtype(arg_LSE, "LSE", 2, 32, 1); + CheckContiguous(arg_LSE, "LSE"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_qk_dim >= -2147483648LL && arg_qk_dim <= 2147483647LL, ValueError) + << "scalar 'qk_dim' value " << arg_qk_dim + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_dim >= -2147483648LL && arg_value_dim <= 2147483647LL, ValueError) + << "scalar 'value_dim' value " << arg_value_dim + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_kv_stride >= -2147483648LL && arg_kv_stride <= 2147483647LL, ValueError) + << "scalar 'kv_stride' value " << arg_kv_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_size >= -2147483648LL && arg_page_size <= 2147483647LL, ValueError) + << "scalar 'page_size' value " << arg_page_size + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_table_width >= -2147483648LL && arg_page_table_width <= 2147483647LL, ValueError) + << "scalar 'page_table_width' value " << arg_page_table_width + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_width >= -2147483648LL && arg_sparse_width <= 2147483647LL, ValueError) + << "scalar 'sparse_width' value " << arg_sparse_width + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_sparse >= -2147483648LL && arg_use_sparse <= 2147483647LL, ValueError) + << "scalar 'use_sparse' value " << arg_use_sparse + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_write_lse >= -2147483648LL && arg_write_lse <= 2147483647LL, ValueError) + << "scalar 'write_lse' value " << arg_write_lse + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_KV, arg_Q, "KV", "Q"); + CheckSameCudaDevice(arg_fp8_lut, arg_Q, "fp8_lut", "Q"); + CheckSameCudaDevice(arg_page_table, arg_Q, "page_table", "Q"); + CheckSameCudaDevice(arg_sparse_indices, arg_Q, "sparse_indices", "Q"); + CheckSameCudaDevice(arg_row_batches, arg_Q, "row_batches", "Q"); + CheckSameCudaDevice(arg_row_seq_lens, arg_Q, "row_seq_lens", "Q"); + CheckSameCudaDevice(arg_O, arg_Q, "O", "Q"); + CheckSameCudaDevice(arg_LSE, arg_Q, "LSE", "Q"); + CheckSameCudaDevice(arg_sinks, arg_Q, "sinks", "Q"); + CheckCurrentCudaDevice(arg_Q, "Q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + void* p_Q = arg_Q.data_ptr(); + void* p_KV = arg_KV.data_ptr(); + void* p_fp8_lut = arg_fp8_lut.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sparse_indices = arg_sparse_indices.data_ptr(); + void* p_row_batches = arg_row_batches.data_ptr(); + void* p_row_seq_lens = arg_row_seq_lens.data_ptr(); + void* p_O = arg_O.data_ptr(); + void* p_LSE = arg_LSE.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_qk_dim = (int32_t)arg_qk_dim; + int32_t v_value_dim = (int32_t)arg_value_dim; + int32_t v_kv_stride = (int32_t)arg_kv_stride; + int32_t v_page_size = (int32_t)arg_page_size; + int32_t v_page_table_width = (int32_t)arg_page_table_width; + int32_t v_sparse_width = (int32_t)arg_sparse_width; + int32_t v_use_sparse = (int32_t)arg_use_sparse; + float v_softmax_scale = (float)arg_softmax_scale; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + int32_t v_write_lse = (int32_t)arg_write_lse; + void* kargs[] = {&p_Q, &p_KV, &p_fp8_lut, &p_page_table, &p_sparse_indices, &p_row_batches, &p_row_seq_lens, &p_O, &p_LSE, &p_sinks, &v_num_heads, &v_qk_dim, &v_value_dim, &v_kv_stride, &v_page_size, &v_page_table_width, &v_sparse_width, &v_use_sparse, &v_softmax_scale, &v_bmm2_scale, &v_enable_sink, &v_write_lse}; + + static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(full_abi_tail_fp8_26b36022eb, "kernel_full_abi_tail_fp8"); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(32u, 1u, 1u); + + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.Launch(kargs, grid, block, stream, 16384u)); +} +} // namespace variant_full_abi_tail_fp8_26b36022eb_f0c2469a092e + +int SelectVariant(int32_t num_heads, int32_t qk_dim, int32_t value_dim, int32_t kv_stride, int32_t page_size, int32_t page_table_width, int32_t sparse_width, int32_t use_sparse, float softmax_scale, float bmm2_scale, int32_t enable_sink, int32_t write_lse, int32_t num_rows) { + if ((num_rows > 0)) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_query, tvm::ffi::TensorView arg_kv, tvm::ffi::TensorView arg_fp8_lut, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sparse_indices, tvm::ffi::TensorView arg_row_batches, tvm::ffi::TensorView arg_row_seq_lens, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_lse, tvm::ffi::TensorView arg_sinks, int32_t num_heads, int32_t qk_dim, int32_t value_dim, int32_t kv_stride, int32_t page_size, int32_t page_table_width, int32_t sparse_width, int32_t use_sparse, float softmax_scale, float bmm2_scale, int32_t enable_sink, int32_t write_lse, int32_t num_rows, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_query, "query"); + CheckCudaTensor(arg_kv, "kv"); + CheckCudaTensor(arg_fp8_lut, "fp8_lut"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sparse_indices, "sparse_indices"); + CheckCudaTensor(arg_row_batches, "row_batches"); + CheckCudaTensor(arg_row_seq_lens, "row_seq_lens"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_lse, "lse"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(num_heads, qk_dim, value_dim, kv_stride, page_size, page_table_width, sparse_width, use_sparse, softmax_scale, bmm2_scale, enable_sink, write_lse, num_rows); + switch (v) { + case 0: { + { + int64_t gx = num_rows; + int64_t gy = num_heads; + int64_t gz = 1; + mla_host_shim::variant_full_abi_tail_fp8_26b36022eb_f0c2469a092e::Run(arg_query, arg_kv, arg_fp8_lut, arg_page_table, arg_sparse_indices, arg_row_batches, arg_row_seq_lens, arg_output, arg_lse, arg_sinks, num_heads, qk_dim, value_dim, kv_stride, page_size, page_table_width, sparse_width, use_sparse, softmax_scale, bmm2_scale, enable_sink, write_lse, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_fp8_tail shape (num_heads=" << num_heads << ", qk_dim=" << qk_dim << ", value_dim=" << value_dim << ", kv_stride=" << kv_stride << ", page_size=" << page_size << ", page_table_width=" << page_table_width << ", sparse_width=" << sparse_width << ", use_sparse=" << use_sparse << ", softmax_scale=" << softmax_scale << ", bmm2_scale=" << bmm2_scale << ", enable_sink=" << enable_sink << ", write_lse=" << write_lse << ", num_rows=" << num_rows << ")"; + } +} + +int64_t RouteOf(int32_t num_heads, int32_t qk_dim, int32_t value_dim, int32_t kv_stride, int32_t page_size, int32_t page_table_width, int32_t sparse_width, int32_t use_sparse, float softmax_scale, float bmm2_scale, int32_t enable_sink, int32_t write_lse, int32_t num_rows) { return SelectVariant(num_heads, qk_dim, value_dim, kv_stride, page_size, page_table_width, sparse_width, use_sparse, softmax_scale, bmm2_scale, enable_sink, write_lse, num_rows); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_clc.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_clc.cpp new file mode 100644 index 00000000000..ed50cd8c617 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_clc.cpp @@ -0,0 +1,2810 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_5f26657c8f); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_5f6f81e2ad); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_823f60e1b7); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_a6a70d7004); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_d003c22422); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_d3892bb3e6); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_d56d2460cb); +TVM_FFI_EMBED_CUBIN(mla_decode_exact_live_bf16_clc_fb6e96f4d8); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +namespace variant_mla_decode_exact_live_bf16_clc_d56d2460cb_ceeecafef56b { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_d56d2460cb::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_d56d2460cb_ceeecafef56b + +namespace variant_mla_decode_exact_live_bf16_clc_d003c22422_d1a0fc38928b { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_d003c22422::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_d003c22422_d1a0fc38928b + +namespace variant_mla_decode_exact_live_bf16_clc_a6a70d7004_92f6ccbf4561 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_a6a70d7004::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_a6a70d7004_92f6ccbf4561 + +namespace variant_mla_decode_exact_live_bf16_clc_fb6e96f4d8_e03a032a0a9a { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_fb6e96f4d8::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_fb6e96f4d8_e03a032a0a9a + +namespace variant_mla_decode_exact_live_bf16_clc_d3892bb3e6_9ae1170423f2 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_d3892bb3e6::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_d3892bb3e6_9ae1170423f2 + +namespace variant_mla_decode_exact_live_bf16_clc_5f26657c8f_13913bd93202 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_5f26657c8f::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_5f26657c8f_13913bd93202 + +namespace variant_mla_decode_exact_live_bf16_clc_5f6f81e2ad_f935d051e5c5 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_5f6f81e2ad::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_5f6f81e2ad_f935d051e5c5 + +namespace variant_mla_decode_exact_live_bf16_clc_823f60e1b7_30db214ceb9a { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 4, 16, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_exact_live_bf16_clc_823f60e1b7::Global()->mod.GetKernel("kernel_mla_decode_exact_live_bf16_clc"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 230912, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 230912u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 230912u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_exact_live_bf16_clc_823f60e1b7_30db214ceb9a + +int SelectVariant(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z, int32_t max_kv_len, int32_t last_kv_len, int32_t source_selector_eligibility) { + if (((((total_work_items == 512) && (value_split_count == 1)) && (grid_x == 2)) && (grid_z == 512))) return 0; + if (((((((total_work_items == 4096) && (value_split_count == 4)) && (max_pages_per_seq == 32)) && (grid_x == 8)) && (grid_z == 1024)) && ((((source_selector_eligibility == 1) || (source_selector_eligibility == 3)) || (source_selector_eligibility == 5)) || (source_selector_eligibility == 7)))) return 1; + if ((((((((((source_selector_eligibility == 6) || (source_selector_eligibility == 7)) && (total_work_items == 6144)) && (value_split_count == 8)) && (max_kv_len == 1024)) && (last_kv_len == 1024)) && (max_pages_per_seq == 32)) && (grid_x == 16)) && (grid_z == 768))) return 2; + if (((((total_work_items == 6144) && (value_split_count == 8)) && (grid_x == 16)) && (grid_z == 768))) return 3; + if (((((((source_selector_eligibility == 2) || (source_selector_eligibility == 3)) || (source_selector_eligibility == 6)) || (source_selector_eligibility == 7)) && (last_kv_len == max_kv_len)) && ((((((total_work_items == 128) && (value_split_count == 2)) && (max_kv_len == 8192)) && (max_pages_per_seq == 256)) && (grid_x == 4)) && (grid_z == 64)))) return 4; + if (((((((source_selector_eligibility == 2) || (source_selector_eligibility == 3)) || (source_selector_eligibility == 6)) || (source_selector_eligibility == 7)) && (last_kv_len == max_kv_len)) && ((((((total_work_items == 8192) && (value_split_count == 16)) && (max_kv_len == 1024)) && (max_pages_per_seq == 32)) && (grid_x == 32)) && (grid_z == 512)))) return 5; + if (((((((source_selector_eligibility == 2) || (source_selector_eligibility == 3)) || (source_selector_eligibility == 6)) || (source_selector_eligibility == 7)) && (last_kv_len == max_kv_len)) && ((((((total_work_items == 1024) && (value_split_count == 2)) && (max_kv_len == 4096)) && (max_pages_per_seq == 128)) && (grid_x == 4)) && (grid_z == 512)))) return 6; + if ((((total_work_items > 0) && (grid_x > 0)) && (grid_z > 0))) return 7; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_q_rows, tvm::ffi::TensorView arg_kv_pages, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sinks, float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z, int32_t max_kv_len, int32_t last_kv_len, int32_t source_selector_eligibility, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_q_rows, "q_rows"); + CheckCudaTensor(arg_kv_pages, "kv_pages"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z, max_kv_len, last_kv_len, source_selector_eligibility); + switch (v) { + case 0: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_d56d2460cb_ceeecafef56b::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 1: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_d003c22422_d1a0fc38928b::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 2: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_a6a70d7004_92f6ccbf4561::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 3: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_fb6e96f4d8_e03a032a0a9a::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 4: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_d3892bb3e6_9ae1170423f2::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 5: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_5f26657c8f_13913bd93202::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 6: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_5f6f81e2ad_f935d051e5c5::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + case 7: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_exact_live_bf16_clc_823f60e1b7_30db214ceb9a::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_bf16_clc shape (softmax_scale_log2=" << softmax_scale_log2 << ", bmm2_scale=" << bmm2_scale << ", total_work_items=" << total_work_items << ", value_split_count=" << value_split_count << ", max_pages_per_seq=" << max_pages_per_seq << ", enable_sink=" << enable_sink << ", grid_x=" << grid_x << ", grid_z=" << grid_z << ", max_kv_len=" << max_kv_len << ", last_kv_len=" << last_kv_len << ", source_selector_eligibility=" << source_selector_eligibility << ")"; + } +} + +int64_t RouteOf(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z, int32_t max_kv_len, int32_t last_kv_len, int32_t source_selector_eligibility) { return SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z, max_kv_len, last_kv_len, source_selector_eligibility); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_native_split8_pdl.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_native_split8_pdl.cpp new file mode 100644 index 00000000000..a4fcf191e84 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_native_split8_pdl.cpp @@ -0,0 +1,767 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(mla_decode_v32_paged_bf16_c80238a3c3); +TVM_FFI_EMBED_CUBIN(mla_decode_v32_reduce_1db619cce9); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +namespace variant_mla_decode_v32_paged_bf16_c80238a3c3_1fef86a30b84 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_nope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_nope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_nope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_nope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_nope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_nope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_nope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_nope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_rope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_rope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_rope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_rope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_rope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_rope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_rope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_rope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k_nope, TensorView arg_tmap_k_rope, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_num_split, int64_t arg_total_work_items, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k_nope, "tmap_k_nope"); + CheckDtype(arg_tmap_k_nope, "tmap_k_nope", 4, 16, 1); + CheckContiguous(arg_tmap_k_nope, "tmap_k_nope"); + CheckCudaTensor(arg_tmap_k_rope, "tmap_k_rope"); + CheckDtype(arg_tmap_k_rope, "tmap_k_rope", 4, 16, 1); + CheckContiguous(arg_tmap_k_rope, "tmap_k_rope"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) + << "scalar 'num_split' value " << arg_num_split + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k_nope, arg_tmap_q, "tmap_k_nope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_k_rope, arg_tmap_q, "tmap_k_rope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_k_nope = EncodeTma_tmap_k_nope(arg_tmap_k_nope); + CUtensorMap p_tmap_k_rope = EncodeTma_tmap_k_rope(arg_tmap_k_rope); + CUtensorMap p_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_O = arg_O.data_ptr(); + void* p_partial_lse = arg_partial_lse.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_num_split = (int32_t)arg_num_split; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k_nope, &p_tmap_k_rope, &p_tmap_v, &p_O, &p_partial_lse, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_num_split, &v_total_work_items, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_v32_paged_bf16_c80238a3c3::Global()->mod.GetKernel("kernel_mla_decode_v32_paged_bf16"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 231424, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 231424u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 231424u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_v32_paged_bf16_c80238a3c3_1fef86a30b84 + +namespace variant_mla_decode_v32_reduce_1db619cce9_4512a48ed07e { + +void Run(TensorView arg_partial_O, TensorView arg_partial_lse, TensorView arg_O, TensorView arg_sinks, int64_t arg_num_q_heads, int64_t arg_num_split, double arg_bmm2_scale, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_partial_O.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); + CheckContiguous(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) + << "scalar 'num_q_heads' value " << arg_num_q_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) + << "scalar 'num_split' value " << arg_num_split + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_partial_lse, arg_partial_O, "partial_lse", "partial_O"); + CheckSameCudaDevice(arg_O, arg_partial_O, "O", "partial_O"); + CheckSameCudaDevice(arg_sinks, arg_partial_O, "sinks", "partial_O"); + CheckCurrentCudaDevice(arg_partial_O, "partial_O"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + void* p_partial_O = arg_partial_O.data_ptr(); + void* p_partial_lse = arg_partial_lse.data_ptr(); + void* p_O = arg_O.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + int32_t v_num_q_heads = (int32_t)arg_num_q_heads; + int32_t v_num_split = (int32_t)arg_num_split; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_partial_O, &p_partial_lse, &p_O, &p_sinks, &v_num_q_heads, &v_num_split, &v_bmm2_scale, &v_enable_sink}; + + static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(mla_decode_v32_reduce_1db619cce9, "kernel_mla_decode_v32_reduce"); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(128u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[1]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 256u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[1]; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 256u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_v32_reduce_1db619cce9_4512a48ed07e + +int SelectVariant(float softmax_scale_log2, float partial_bmm2_scale, float bmm2_scale, int32_t num_split, int32_t total_work_items, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t num_rows) { + if (((((num_split == 8) && (total_work_items == (num_rows * num_split))) && (num_rows > 0)) && (grid_x > 0))) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_q_rows, tvm::ffi::TensorView arg_kv_half_pages, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sinks, tvm::ffi::TensorView arg_ws_partial_output, tvm::ffi::TensorView arg_ws_partial_lse, float softmax_scale_log2, float partial_bmm2_scale, float bmm2_scale, int32_t num_split, int32_t total_work_items, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t num_rows, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_q_rows, "q_rows"); + CheckCudaTensor(arg_kv_half_pages, "kv_half_pages"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(softmax_scale_log2, partial_bmm2_scale, bmm2_scale, num_split, total_work_items, max_pages_per_seq, enable_sink, grid_x, num_rows); + switch (v) { + case 0: { + CheckCudaTensor(arg_ws_partial_output, "partial_output"); + CheckDtype(arg_ws_partial_output, "partial_output", 4, 16, 1); + TVM_FFI_CHECK(arg_ws_partial_output.numel() == (int64_t)((num_rows) * (128) * (num_split) * (512)), ValueError) + << "workspace 'partial_output' numel " << arg_ws_partial_output.numel() + << " != expected " << (int64_t)((num_rows) * (128) * (num_split) * (512)); + CheckCudaTensor(arg_ws_partial_lse, "partial_lse"); + CheckDtype(arg_ws_partial_lse, "partial_lse", 2, 32, 1); + TVM_FFI_CHECK(arg_ws_partial_lse.numel() == (int64_t)((num_rows) * (128) * (num_split)), ValueError) + << "workspace 'partial_lse' numel " << arg_ws_partial_lse.numel() + << " != expected " << (int64_t)((num_rows) * (128) * (num_split)); + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = 1; + mla_host_shim::variant_mla_decode_v32_paged_bf16_c80238a3c3_1fef86a30b84::Run(arg_q_rows, arg_kv_half_pages, arg_kv_half_pages, arg_kv_half_pages, arg_ws_partial_output, arg_ws_partial_lse, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, partial_bmm2_scale, num_split, total_work_items, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + { + int64_t gx = num_rows; + int64_t gy = 128; + int64_t gz = 1; + mla_host_shim::variant_mla_decode_v32_reduce_1db619cce9_4512a48ed07e::Run(arg_ws_partial_output, arg_ws_partial_lse, arg_output, arg_sinks, 128, num_split, bmm2_scale, enable_sink, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_bf16_native_split8_pdl shape (softmax_scale_log2=" << softmax_scale_log2 << ", partial_bmm2_scale=" << partial_bmm2_scale << ", bmm2_scale=" << bmm2_scale << ", num_split=" << num_split << ", total_work_items=" << total_work_items << ", max_pages_per_seq=" << max_pages_per_seq << ", enable_sink=" << enable_sink << ", grid_x=" << grid_x << ", num_rows=" << num_rows << ")"; + } +} + +int64_t RouteOf(float softmax_scale_log2, float partial_bmm2_scale, float bmm2_scale, int32_t num_split, int32_t total_work_items, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t num_rows) { return SelectVariant(softmax_scale_log2, partial_bmm2_scale, bmm2_scale, num_split, total_work_items, max_pages_per_seq, enable_sink, grid_x, num_rows); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_tail.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_tail.cpp new file mode 100644 index 00000000000..ea53af829b9 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_tail.cpp @@ -0,0 +1,409 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(full_abi_tail_bf16_83451b7317); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +namespace variant_full_abi_tail_bf16_83451b7317_d1d6ab272ce6 { + +void Run(TensorView arg_Q, TensorView arg_KV, TensorView arg_fp8_lut, TensorView arg_page_table, TensorView arg_sparse_indices, TensorView arg_row_batches, TensorView arg_row_seq_lens, TensorView arg_O, TensorView arg_LSE, TensorView arg_sinks, int64_t arg_num_heads, int64_t arg_qk_dim, int64_t arg_value_dim, int64_t arg_kv_stride, int64_t arg_page_size, int64_t arg_page_table_width, int64_t arg_sparse_width, int64_t arg_use_sparse, double arg_softmax_scale, double arg_bmm2_scale, int64_t arg_enable_sink, int64_t arg_write_lse, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_Q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_Q, "Q"); + CheckDtype(arg_Q, "Q", 4, 16, 1); + CheckContiguous(arg_Q, "Q"); + CheckCudaTensor(arg_KV, "KV"); + CheckDtype(arg_KV, "KV", 4, 16, 1); + CheckContiguous(arg_KV, "KV"); + CheckCudaTensor(arg_fp8_lut, "fp8_lut"); + CheckDtype(arg_fp8_lut, "fp8_lut", 2, 32, 1); + CheckContiguous(arg_fp8_lut, "fp8_lut"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sparse_indices, "sparse_indices"); + CheckDtype(arg_sparse_indices, "sparse_indices", 0, 32, 1); + CheckContiguous(arg_sparse_indices, "sparse_indices"); + CheckCudaTensor(arg_row_batches, "row_batches"); + CheckDtype(arg_row_batches, "row_batches", 0, 32, 1); + CheckContiguous(arg_row_batches, "row_batches"); + CheckCudaTensor(arg_row_seq_lens, "row_seq_lens"); + CheckDtype(arg_row_seq_lens, "row_seq_lens", 0, 32, 1); + CheckContiguous(arg_row_seq_lens, "row_seq_lens"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_LSE, "LSE"); + CheckDtype(arg_LSE, "LSE", 2, 32, 1); + CheckContiguous(arg_LSE, "LSE"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_qk_dim >= -2147483648LL && arg_qk_dim <= 2147483647LL, ValueError) + << "scalar 'qk_dim' value " << arg_qk_dim + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_dim >= -2147483648LL && arg_value_dim <= 2147483647LL, ValueError) + << "scalar 'value_dim' value " << arg_value_dim + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_kv_stride >= -2147483648LL && arg_kv_stride <= 2147483647LL, ValueError) + << "scalar 'kv_stride' value " << arg_kv_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_size >= -2147483648LL && arg_page_size <= 2147483647LL, ValueError) + << "scalar 'page_size' value " << arg_page_size + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_table_width >= -2147483648LL && arg_page_table_width <= 2147483647LL, ValueError) + << "scalar 'page_table_width' value " << arg_page_table_width + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_width >= -2147483648LL && arg_sparse_width <= 2147483647LL, ValueError) + << "scalar 'sparse_width' value " << arg_sparse_width + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_sparse >= -2147483648LL && arg_use_sparse <= 2147483647LL, ValueError) + << "scalar 'use_sparse' value " << arg_use_sparse + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_write_lse >= -2147483648LL && arg_write_lse <= 2147483647LL, ValueError) + << "scalar 'write_lse' value " << arg_write_lse + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_KV, arg_Q, "KV", "Q"); + CheckSameCudaDevice(arg_fp8_lut, arg_Q, "fp8_lut", "Q"); + CheckSameCudaDevice(arg_page_table, arg_Q, "page_table", "Q"); + CheckSameCudaDevice(arg_sparse_indices, arg_Q, "sparse_indices", "Q"); + CheckSameCudaDevice(arg_row_batches, arg_Q, "row_batches", "Q"); + CheckSameCudaDevice(arg_row_seq_lens, arg_Q, "row_seq_lens", "Q"); + CheckSameCudaDevice(arg_O, arg_Q, "O", "Q"); + CheckSameCudaDevice(arg_LSE, arg_Q, "LSE", "Q"); + CheckSameCudaDevice(arg_sinks, arg_Q, "sinks", "Q"); + CheckCurrentCudaDevice(arg_Q, "Q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + void* p_Q = arg_Q.data_ptr(); + void* p_KV = arg_KV.data_ptr(); + void* p_fp8_lut = arg_fp8_lut.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sparse_indices = arg_sparse_indices.data_ptr(); + void* p_row_batches = arg_row_batches.data_ptr(); + void* p_row_seq_lens = arg_row_seq_lens.data_ptr(); + void* p_O = arg_O.data_ptr(); + void* p_LSE = arg_LSE.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_qk_dim = (int32_t)arg_qk_dim; + int32_t v_value_dim = (int32_t)arg_value_dim; + int32_t v_kv_stride = (int32_t)arg_kv_stride; + int32_t v_page_size = (int32_t)arg_page_size; + int32_t v_page_table_width = (int32_t)arg_page_table_width; + int32_t v_sparse_width = (int32_t)arg_sparse_width; + int32_t v_use_sparse = (int32_t)arg_use_sparse; + float v_softmax_scale = (float)arg_softmax_scale; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + int32_t v_write_lse = (int32_t)arg_write_lse; + void* kargs[] = {&p_Q, &p_KV, &p_fp8_lut, &p_page_table, &p_sparse_indices, &p_row_batches, &p_row_seq_lens, &p_O, &p_LSE, &p_sinks, &v_num_heads, &v_qk_dim, &v_value_dim, &v_kv_stride, &v_page_size, &v_page_table_width, &v_sparse_width, &v_use_sparse, &v_softmax_scale, &v_bmm2_scale, &v_enable_sink, &v_write_lse}; + + static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(full_abi_tail_bf16_83451b7317, "kernel_full_abi_tail_bf16"); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(32u, 1u, 1u); + + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.Launch(kargs, grid, block, stream, 16384u)); +} +} // namespace variant_full_abi_tail_bf16_83451b7317_d1d6ab272ce6 + +int SelectVariant(int32_t num_heads, int32_t qk_dim, int32_t value_dim, int32_t kv_stride, int32_t page_size, int32_t page_table_width, int32_t sparse_width, int32_t use_sparse, float softmax_scale, float bmm2_scale, int32_t enable_sink, int32_t write_lse, int32_t num_rows) { + if ((num_rows > 0)) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_query, tvm::ffi::TensorView arg_kv, tvm::ffi::TensorView arg_fp8_lut, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sparse_indices, tvm::ffi::TensorView arg_row_batches, tvm::ffi::TensorView arg_row_seq_lens, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_lse, tvm::ffi::TensorView arg_sinks, int32_t num_heads, int32_t qk_dim, int32_t value_dim, int32_t kv_stride, int32_t page_size, int32_t page_table_width, int32_t sparse_width, int32_t use_sparse, float softmax_scale, float bmm2_scale, int32_t enable_sink, int32_t write_lse, int32_t num_rows, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_query, "query"); + CheckCudaTensor(arg_kv, "kv"); + CheckCudaTensor(arg_fp8_lut, "fp8_lut"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sparse_indices, "sparse_indices"); + CheckCudaTensor(arg_row_batches, "row_batches"); + CheckCudaTensor(arg_row_seq_lens, "row_seq_lens"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_lse, "lse"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(num_heads, qk_dim, value_dim, kv_stride, page_size, page_table_width, sparse_width, use_sparse, softmax_scale, bmm2_scale, enable_sink, write_lse, num_rows); + switch (v) { + case 0: { + { + int64_t gx = num_rows; + int64_t gy = num_heads; + int64_t gz = 1; + mla_host_shim::variant_full_abi_tail_bf16_83451b7317_d1d6ab272ce6::Run(arg_query, arg_kv, arg_fp8_lut, arg_page_table, arg_sparse_indices, arg_row_batches, arg_row_seq_lens, arg_output, arg_lse, arg_sinks, num_heads, qk_dim, value_dim, kv_stride, page_size, page_table_width, sparse_width, use_sparse, softmax_scale, bmm2_scale, enable_sink, write_lse, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_bf16_tail shape (num_heads=" << num_heads << ", qk_dim=" << qk_dim << ", value_dim=" << value_dim << ", kv_stride=" << kv_stride << ", page_size=" << page_size << ", page_table_width=" << page_table_width << ", sparse_width=" << sparse_width << ", use_sparse=" << use_sparse << ", softmax_scale=" << softmax_scale << ", bmm2_scale=" << bmm2_scale << ", enable_sink=" << enable_sink << ", write_lse=" << write_lse << ", num_rows=" << num_rows << ")"; + } +} + +int64_t RouteOf(int32_t num_heads, int32_t qk_dim, int32_t value_dim, int32_t kv_stride, int32_t page_size, int32_t page_table_width, int32_t sparse_width, int32_t use_sparse, float softmax_scale, float bmm2_scale, int32_t enable_sink, int32_t write_lse, int32_t num_rows) { return SelectVariant(num_heads, qk_dim, value_dim, kv_stride, page_size, page_table_width, sparse_width, use_sparse, softmax_scale, bmm2_scale, enable_sink, write_lse, num_rows); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_unsplit.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_unsplit.cpp new file mode 100644 index 00000000000..e2a3bc16241 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_unsplit.cpp @@ -0,0 +1,768 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(mla_decode_value_split_paged_bf16_5c36fe8228); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +struct TmaDeviceArena { + static constexpr size_t kSlotsPerChunk = 256; + static constexpr size_t kMaxSlots = 4096; + std::vector chunks; + size_t used = 0; +}; + +// Immutable, process-lifetime device tensor-map slots for the pointer ABI. +// A slot is never rewritten: different descriptor bytes always get a new +// address, so concurrent streams cannot observe a partially updated map. The +// chunked arena caps storage at 512 KiB per CUDA context in this host module. +static inline void* TmaDeviceSlot( + const CUtensorMap& tm, + int device_id, + cudaStream_t stream) { + static std::mutex mu; + static auto* slots = new std::unordered_map(); + static auto* arenas = new std::unordered_map(); + + // Device allocations are context-owned. Resolve and validate the active + // context before cache lookup so a warm entry can never bypass the same + // checks as a cold entry or leak a pointer across contexts on one device. + CUcontext current_context = nullptr; + CUresult result = cuCtxGetCurrent(¤t_context); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_context != nullptr, RuntimeError) + << "pointer TMA ABI requires an active CUDA context: CUresult=" + << static_cast(result); + CUdevice current_device = -1; + result = cuCtxGetDevice(¤t_device); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_device == device_id, RuntimeError) + << "TMA descriptor device mismatch: current=" << current_device + << ", tensor=" << device_id; + + std::string key = + std::to_string(reinterpret_cast(current_context)); + key.push_back(':'); + key.append(reinterpret_cast(&tm), sizeof(CUtensorMap)); + std::lock_guard lock(mu); + + // CUcontext is an opaque, recyclable handle. Before trusting a warm arena, + // prove that its first allocation still belongs to this live context. A + // destroyed context makes the pointer query fail; a subsequently reused + // handle must therefore discard the stale pointer and descriptor entries. + auto arena_it = arenas->find(current_context); + if (arena_it != arenas->end() && !arena_it->second.chunks.empty()) { + CUcontext allocation_context = nullptr; + result = cuPointerGetAttribute( + &allocation_context, + CU_POINTER_ATTRIBUTE_CONTEXT, + arena_it->second.chunks.front()); + if (result != CUDA_SUCCESS || allocation_context != current_context) { + const std::string prefix = + std::to_string(reinterpret_cast(current_context)) + ":"; + for (auto slot_it = slots->begin(); slot_it != slots->end();) { + if (slot_it->first.compare(0, prefix.size(), prefix) == 0) { + slot_it = slots->erase(slot_it); + } else { + ++slot_it; + } + } + arenas->erase(arena_it); + } + } + + auto it = slots->find(key); + if (it != slots->end()) return it->second; + + CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; + result = cuStreamIsCapturing( + reinterpret_cast(stream), &capture_status); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuStreamIsCapturing for TMA descriptor slot failed: CUresult=" + << static_cast(result); + TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) + << "pointer TMA ABI cannot create a new device descriptor slot inside " + "CUDA Graph capture; prewarm this exact tensor/layout binding or " + "compile with tma_abi='grid_constant'"; + + TmaDeviceArena& arena = (*arenas)[current_context]; + TVM_FFI_CHECK(arena.used < TmaDeviceArena::kMaxSlots, RuntimeError) + << "pointer TMA ABI exhausted its immutable descriptor arena in CUDA " + "context " << current_context << " on device " << device_id + << " (capacity=" << TmaDeviceArena::kMaxSlots + << "); reuse tensor/layout bindings or compile with tma_abi='grid_constant'"; + if (arena.used % TmaDeviceArena::kSlotsPerChunk == 0) { + CUdeviceptr chunk = 0; + result = cuMemAlloc( + &chunk, + TmaDeviceArena::kSlotsPerChunk * sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemAlloc for TMA descriptor arena failed: CUresult=" + << static_cast(result); + arena.chunks.push_back(chunk); + } + size_t chunk_index = arena.used / TmaDeviceArena::kSlotsPerChunk; + size_t slot_index = arena.used % TmaDeviceArena::kSlotsPerChunk; + CUdeviceptr dev = arena.chunks[chunk_index] + + slot_index * sizeof(CUtensorMap); + result = cuMemcpyHtoD(dev, &tm, sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemcpyHtoD for TMA descriptor slot failed: CUresult=" + << static_cast(result); + ++arena.used; + void* pointer = reinterpret_cast(static_cast(dev)); + (*slots)[key] = pointer; + return pointer; +} + +namespace variant_mla_decode_value_split_paged_bf16_5c36fe8228_a5bebae701d8 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_nope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_nope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_nope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_nope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_nope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_nope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_nope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_nope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_rope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_rope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_rope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_rope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_rope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_rope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_rope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_rope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k_nope, TensorView arg_tmap_k_rope, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k_nope, "tmap_k_nope"); + CheckDtype(arg_tmap_k_nope, "tmap_k_nope", 4, 16, 1); + CheckContiguous(arg_tmap_k_nope, "tmap_k_nope"); + CheckCudaTensor(arg_tmap_k_rope, "tmap_k_rope"); + CheckDtype(arg_tmap_k_rope, "tmap_k_rope", 4, 16, 1); + CheckContiguous(arg_tmap_k_rope, "tmap_k_rope"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k_nope, arg_tmap_q, "tmap_k_nope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_k_rope, arg_tmap_q, "tmap_k_rope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + void* p_tmap_q = TmaDeviceSlot(h_tmap_q, arg_tmap_q.device().device_id, stream); + CUtensorMap h_tmap_k_nope = EncodeTma_tmap_k_nope(arg_tmap_k_nope); + void* p_tmap_k_nope = TmaDeviceSlot(h_tmap_k_nope, arg_tmap_k_nope.device().device_id, stream); + CUtensorMap h_tmap_k_rope = EncodeTma_tmap_k_rope(arg_tmap_k_rope); + void* p_tmap_k_rope = TmaDeviceSlot(h_tmap_k_rope, arg_tmap_k_rope.device().device_id, stream); + CUtensorMap h_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_tmap_v = TmaDeviceSlot(h_tmap_v, arg_tmap_v.device().device_id, stream); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k_nope, &p_tmap_k_rope, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_value_split_paged_bf16_5c36fe8228::Global()->mod.GetKernel("kernel_mla_decode_value_split_paged_bf16"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 231936, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[2]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 231936u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[2]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 231936u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_value_split_paged_bf16_5c36fe8228_a5bebae701d8 + +int SelectVariant(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z) { + if ((((total_work_items > 0) && (grid_x > 0)) && (grid_z > 0))) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_q_rows, tvm::ffi::TensorView arg_kv_pages, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sinks, float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_q_rows, "q_rows"); + CheckCudaTensor(arg_kv_pages, "kv_pages"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z); + switch (v) { + case 0: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_value_split_paged_bf16_5c36fe8228_a5bebae701d8::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_bf16_unsplit shape (softmax_scale_log2=" << softmax_scale_log2 << ", bmm2_scale=" << bmm2_scale << ", total_work_items=" << total_work_items << ", value_split_count=" << value_split_count << ", max_pages_per_seq=" << max_pages_per_seq << ", enable_sink=" << enable_sink << ", grid_x=" << grid_x << ", grid_z=" << grid_z << ")"; + } +} + +int64_t RouteOf(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z) { return SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_vhalf.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_vhalf.cpp new file mode 100644 index 00000000000..b3b51c6aea6 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_vhalf.cpp @@ -0,0 +1,768 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(mla_decode_value_split_paged_bf16_fa7b59f3a7); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +struct TmaDeviceArena { + static constexpr size_t kSlotsPerChunk = 256; + static constexpr size_t kMaxSlots = 4096; + std::vector chunks; + size_t used = 0; +}; + +// Immutable, process-lifetime device tensor-map slots for the pointer ABI. +// A slot is never rewritten: different descriptor bytes always get a new +// address, so concurrent streams cannot observe a partially updated map. The +// chunked arena caps storage at 512 KiB per CUDA context in this host module. +static inline void* TmaDeviceSlot( + const CUtensorMap& tm, + int device_id, + cudaStream_t stream) { + static std::mutex mu; + static auto* slots = new std::unordered_map(); + static auto* arenas = new std::unordered_map(); + + // Device allocations are context-owned. Resolve and validate the active + // context before cache lookup so a warm entry can never bypass the same + // checks as a cold entry or leak a pointer across contexts on one device. + CUcontext current_context = nullptr; + CUresult result = cuCtxGetCurrent(¤t_context); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_context != nullptr, RuntimeError) + << "pointer TMA ABI requires an active CUDA context: CUresult=" + << static_cast(result); + CUdevice current_device = -1; + result = cuCtxGetDevice(¤t_device); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_device == device_id, RuntimeError) + << "TMA descriptor device mismatch: current=" << current_device + << ", tensor=" << device_id; + + std::string key = + std::to_string(reinterpret_cast(current_context)); + key.push_back(':'); + key.append(reinterpret_cast(&tm), sizeof(CUtensorMap)); + std::lock_guard lock(mu); + + // CUcontext is an opaque, recyclable handle. Before trusting a warm arena, + // prove that its first allocation still belongs to this live context. A + // destroyed context makes the pointer query fail; a subsequently reused + // handle must therefore discard the stale pointer and descriptor entries. + auto arena_it = arenas->find(current_context); + if (arena_it != arenas->end() && !arena_it->second.chunks.empty()) { + CUcontext allocation_context = nullptr; + result = cuPointerGetAttribute( + &allocation_context, + CU_POINTER_ATTRIBUTE_CONTEXT, + arena_it->second.chunks.front()); + if (result != CUDA_SUCCESS || allocation_context != current_context) { + const std::string prefix = + std::to_string(reinterpret_cast(current_context)) + ":"; + for (auto slot_it = slots->begin(); slot_it != slots->end();) { + if (slot_it->first.compare(0, prefix.size(), prefix) == 0) { + slot_it = slots->erase(slot_it); + } else { + ++slot_it; + } + } + arenas->erase(arena_it); + } + } + + auto it = slots->find(key); + if (it != slots->end()) return it->second; + + CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; + result = cuStreamIsCapturing( + reinterpret_cast(stream), &capture_status); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuStreamIsCapturing for TMA descriptor slot failed: CUresult=" + << static_cast(result); + TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) + << "pointer TMA ABI cannot create a new device descriptor slot inside " + "CUDA Graph capture; prewarm this exact tensor/layout binding or " + "compile with tma_abi='grid_constant'"; + + TmaDeviceArena& arena = (*arenas)[current_context]; + TVM_FFI_CHECK(arena.used < TmaDeviceArena::kMaxSlots, RuntimeError) + << "pointer TMA ABI exhausted its immutable descriptor arena in CUDA " + "context " << current_context << " on device " << device_id + << " (capacity=" << TmaDeviceArena::kMaxSlots + << "); reuse tensor/layout bindings or compile with tma_abi='grid_constant'"; + if (arena.used % TmaDeviceArena::kSlotsPerChunk == 0) { + CUdeviceptr chunk = 0; + result = cuMemAlloc( + &chunk, + TmaDeviceArena::kSlotsPerChunk * sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemAlloc for TMA descriptor arena failed: CUresult=" + << static_cast(result); + arena.chunks.push_back(chunk); + } + size_t chunk_index = arena.used / TmaDeviceArena::kSlotsPerChunk; + size_t slot_index = arena.used % TmaDeviceArena::kSlotsPerChunk; + CUdeviceptr dev = arena.chunks[chunk_index] + + slot_index * sizeof(CUtensorMap); + result = cuMemcpyHtoD(dev, &tm, sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemcpyHtoD for TMA descriptor slot failed: CUresult=" + << static_cast(result); + ++arena.used; + void* pointer = reinterpret_cast(static_cast(dev)); + (*slots)[key] = pointer; + return pointer; +} + +namespace variant_mla_decode_value_split_paged_bf16_fa7b59f3a7_ed851bc7fbe4 { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_nope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_nope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_nope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_nope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_nope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_nope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_nope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_nope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_rope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_rope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_rope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_rope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_rope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_rope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_rope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_rope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k_nope, TensorView arg_tmap_k_rope, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k_nope, "tmap_k_nope"); + CheckDtype(arg_tmap_k_nope, "tmap_k_nope", 4, 16, 1); + CheckContiguous(arg_tmap_k_nope, "tmap_k_nope"); + CheckCudaTensor(arg_tmap_k_rope, "tmap_k_rope"); + CheckDtype(arg_tmap_k_rope, "tmap_k_rope", 4, 16, 1); + CheckContiguous(arg_tmap_k_rope, "tmap_k_rope"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k_nope, arg_tmap_q, "tmap_k_nope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_k_rope, arg_tmap_q, "tmap_k_rope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + void* p_tmap_q = TmaDeviceSlot(h_tmap_q, arg_tmap_q.device().device_id, stream); + CUtensorMap h_tmap_k_nope = EncodeTma_tmap_k_nope(arg_tmap_k_nope); + void* p_tmap_k_nope = TmaDeviceSlot(h_tmap_k_nope, arg_tmap_k_nope.device().device_id, stream); + CUtensorMap h_tmap_k_rope = EncodeTma_tmap_k_rope(arg_tmap_k_rope); + void* p_tmap_k_rope = TmaDeviceSlot(h_tmap_k_rope, arg_tmap_k_rope.device().device_id, stream); + CUtensorMap h_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_tmap_v = TmaDeviceSlot(h_tmap_v, arg_tmap_v.device().device_id, stream); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k_nope, &p_tmap_k_rope, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_value_split_paged_bf16_fa7b59f3a7::Global()->mod.GetKernel("kernel_mla_decode_value_split_paged_bf16"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 231424, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[2]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 231424u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[2]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 231424u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_value_split_paged_bf16_fa7b59f3a7_ed851bc7fbe4 + +int SelectVariant(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z) { + if ((((total_work_items > 0) && (grid_x > 0)) && (grid_z > 0))) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_q_rows, tvm::ffi::TensorView arg_kv_pages, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sinks, float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_q_rows, "q_rows"); + CheckCudaTensor(arg_kv_pages, "kv_pages"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z); + switch (v) { + case 0: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_value_split_paged_bf16_fa7b59f3a7_ed851bc7fbe4::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_bf16_vhalf shape (softmax_scale_log2=" << softmax_scale_log2 << ", bmm2_scale=" << bmm2_scale << ", total_work_items=" << total_work_items << ", value_split_count=" << value_split_count << ", max_pages_per_seq=" << max_pages_per_seq << ", enable_sink=" << enable_sink << ", grid_x=" << grid_x << ", grid_z=" << grid_z << ")"; + } +} + +int64_t RouteOf(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z) { return SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_vquarter.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_vquarter.cpp new file mode 100644 index 00000000000..bf2e5f73ed2 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_bf16_vquarter.cpp @@ -0,0 +1,768 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(mla_decode_value_split_paged_bf16_ee41d3ceea); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +struct TmaDeviceArena { + static constexpr size_t kSlotsPerChunk = 256; + static constexpr size_t kMaxSlots = 4096; + std::vector chunks; + size_t used = 0; +}; + +// Immutable, process-lifetime device tensor-map slots for the pointer ABI. +// A slot is never rewritten: different descriptor bytes always get a new +// address, so concurrent streams cannot observe a partially updated map. The +// chunked arena caps storage at 512 KiB per CUDA context in this host module. +static inline void* TmaDeviceSlot( + const CUtensorMap& tm, + int device_id, + cudaStream_t stream) { + static std::mutex mu; + static auto* slots = new std::unordered_map(); + static auto* arenas = new std::unordered_map(); + + // Device allocations are context-owned. Resolve and validate the active + // context before cache lookup so a warm entry can never bypass the same + // checks as a cold entry or leak a pointer across contexts on one device. + CUcontext current_context = nullptr; + CUresult result = cuCtxGetCurrent(¤t_context); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_context != nullptr, RuntimeError) + << "pointer TMA ABI requires an active CUDA context: CUresult=" + << static_cast(result); + CUdevice current_device = -1; + result = cuCtxGetDevice(¤t_device); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_device == device_id, RuntimeError) + << "TMA descriptor device mismatch: current=" << current_device + << ", tensor=" << device_id; + + std::string key = + std::to_string(reinterpret_cast(current_context)); + key.push_back(':'); + key.append(reinterpret_cast(&tm), sizeof(CUtensorMap)); + std::lock_guard lock(mu); + + // CUcontext is an opaque, recyclable handle. Before trusting a warm arena, + // prove that its first allocation still belongs to this live context. A + // destroyed context makes the pointer query fail; a subsequently reused + // handle must therefore discard the stale pointer and descriptor entries. + auto arena_it = arenas->find(current_context); + if (arena_it != arenas->end() && !arena_it->second.chunks.empty()) { + CUcontext allocation_context = nullptr; + result = cuPointerGetAttribute( + &allocation_context, + CU_POINTER_ATTRIBUTE_CONTEXT, + arena_it->second.chunks.front()); + if (result != CUDA_SUCCESS || allocation_context != current_context) { + const std::string prefix = + std::to_string(reinterpret_cast(current_context)) + ":"; + for (auto slot_it = slots->begin(); slot_it != slots->end();) { + if (slot_it->first.compare(0, prefix.size(), prefix) == 0) { + slot_it = slots->erase(slot_it); + } else { + ++slot_it; + } + } + arenas->erase(arena_it); + } + } + + auto it = slots->find(key); + if (it != slots->end()) return it->second; + + CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; + result = cuStreamIsCapturing( + reinterpret_cast(stream), &capture_status); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuStreamIsCapturing for TMA descriptor slot failed: CUresult=" + << static_cast(result); + TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) + << "pointer TMA ABI cannot create a new device descriptor slot inside " + "CUDA Graph capture; prewarm this exact tensor/layout binding or " + "compile with tma_abi='grid_constant'"; + + TmaDeviceArena& arena = (*arenas)[current_context]; + TVM_FFI_CHECK(arena.used < TmaDeviceArena::kMaxSlots, RuntimeError) + << "pointer TMA ABI exhausted its immutable descriptor arena in CUDA " + "context " << current_context << " on device " << device_id + << " (capacity=" << TmaDeviceArena::kMaxSlots + << "); reuse tensor/layout bindings or compile with tma_abi='grid_constant'"; + if (arena.used % TmaDeviceArena::kSlotsPerChunk == 0) { + CUdeviceptr chunk = 0; + result = cuMemAlloc( + &chunk, + TmaDeviceArena::kSlotsPerChunk * sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemAlloc for TMA descriptor arena failed: CUresult=" + << static_cast(result); + arena.chunks.push_back(chunk); + } + size_t chunk_index = arena.used / TmaDeviceArena::kSlotsPerChunk; + size_t slot_index = arena.used % TmaDeviceArena::kSlotsPerChunk; + CUdeviceptr dev = arena.chunks[chunk_index] + + slot_index * sizeof(CUtensorMap); + result = cuMemcpyHtoD(dev, &tm, sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemcpyHtoD for TMA descriptor slot failed: CUresult=" + << static_cast(result); + ++arena.used; + void* pointer = reinterpret_cast(static_cast(dev)); + (*slots)[key] = pointer; + return pointer; +} + +namespace variant_mla_decode_value_split_paged_bf16_ee41d3ceea_3d9d1dc71fbc { + +// 3D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_q' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_q"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_q' physical strides must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[3] = {(uint64_t)(64), (uint64_t)(outer1), (uint64_t)((d1 / 64))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) + << "TMA box (64, 64, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + uint64_t global_strides[2] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + }; + uint32_t box_dim[3] = {64u, 64u, 1u}; + uint32_t elem_strides[3] = {1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (3D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_nope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_nope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_nope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_nope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_nope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_nope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_nope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_nope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_nope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_k_rope' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k_rope(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_k_rope' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k_rope' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_k_rope' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_k_rope' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_k_rope'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k_rope' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_k_rope') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + TVM_FFI_CHECK(d1 > 0 && d2 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer2 = t.numel() / (d1 * d2); + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_v' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(outer2)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 32u <= global_dim[1] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 32, 1, 1) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 32u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k_nope, TensorView arg_tmap_k_rope, TensorView arg_tmap_v, TensorView arg_O, TensorView arg_seq_lens_kv, TensorView arg_page_table, TensorView arg_sinks, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_total_work_items, int64_t arg_value_split_count, int64_t arg_max_pages_per_seq, int64_t arg_enable_sink, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckCudaTensor(arg_tmap_k_nope, "tmap_k_nope"); + CheckDtype(arg_tmap_k_nope, "tmap_k_nope", 4, 16, 1); + CheckContiguous(arg_tmap_k_nope, "tmap_k_nope"); + CheckCudaTensor(arg_tmap_k_rope, "tmap_k_rope"); + CheckDtype(arg_tmap_k_rope, "tmap_k_rope", 4, 16, 1); + CheckContiguous(arg_tmap_k_rope, "tmap_k_rope"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 4, 16, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_split_count >= -2147483648LL && arg_value_split_count <= 2147483647LL, ValueError) + << "scalar 'value_split_count' value " << arg_value_split_count + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k_nope, arg_tmap_q, "tmap_k_nope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_k_rope, arg_tmap_q, "tmap_k_rope", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + void* p_tmap_q = TmaDeviceSlot(h_tmap_q, arg_tmap_q.device().device_id, stream); + CUtensorMap h_tmap_k_nope = EncodeTma_tmap_k_nope(arg_tmap_k_nope); + void* p_tmap_k_nope = TmaDeviceSlot(h_tmap_k_nope, arg_tmap_k_nope.device().device_id, stream); + CUtensorMap h_tmap_k_rope = EncodeTma_tmap_k_rope(arg_tmap_k_rope); + void* p_tmap_k_rope = TmaDeviceSlot(h_tmap_k_rope, arg_tmap_k_rope.device().device_id, stream); + CUtensorMap h_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_tmap_v = TmaDeviceSlot(h_tmap_v, arg_tmap_v.device().device_id, stream); + void* p_O = arg_O.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + int32_t v_value_split_count = (int32_t)arg_value_split_count; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + void* kargs[] = {&p_tmap_q, &p_tmap_k_nope, &p_tmap_k_rope, &p_tmap_v, &p_O, &p_seq_lens_kv, &p_page_table, &p_sinks, &v_softmax_scale_log2, &v_bmm2_scale, &v_total_work_items, &v_value_split_count, &v_max_pages_per_seq, &v_enable_sink}; + + static auto kernel = EmbedCubinModule_mla_decode_value_split_paged_bf16_ee41d3ceea::Global()->mod.GetKernel("kernel_mla_decode_value_split_paged_bf16"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 231424, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[2]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 231424u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[2]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 231424u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_mla_decode_value_split_paged_bf16_ee41d3ceea_3d9d1dc71fbc + +int SelectVariant(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z) { + if ((((total_work_items > 0) && (grid_x > 0)) && (grid_z > 0))) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_q_rows, tvm::ffi::TensorView arg_kv_pages, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sinks, float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_q_rows, "q_rows"); + CheckCudaTensor(arg_kv_pages, "kv_pages"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z); + switch (v) { + case 0: { + { + int64_t gx = grid_x; + int64_t gy = 1; + int64_t gz = grid_z; + mla_host_shim::variant_mla_decode_value_split_paged_bf16_ee41d3ceea_3d9d1dc71fbc::Run(arg_q_rows, arg_kv_pages, arg_kv_pages, arg_kv_pages, arg_output, arg_seq_lens, arg_page_table, arg_sinks, softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_bf16_vquarter shape (softmax_scale_log2=" << softmax_scale_log2 << ", bmm2_scale=" << bmm2_scale << ", total_work_items=" << total_work_items << ", value_split_count=" << value_split_count << ", max_pages_per_seq=" << max_pages_per_seq << ", enable_sink=" << enable_sink << ", grid_x=" << grid_x << ", grid_z=" << grid_z << ")"; + } +} + +int64_t RouteOf(float softmax_scale_log2, float bmm2_scale, int32_t total_work_items, int32_t value_split_count, int32_t max_pages_per_seq, int32_t enable_sink, int32_t grid_x, int32_t grid_z) { return SelectVariant(softmax_scale_log2, bmm2_scale, total_work_items, value_split_count, max_pages_per_seq, enable_sink, grid_x, grid_z); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_fp8_p32_qk_l2.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_fp8_p32_qk_l2.cpp new file mode 100644 index 00000000000..cf635d85c99 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_fp8_p32_qk_l2.cpp @@ -0,0 +1,560 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(fp8_p32_exact_producer_105fd84009); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +namespace variant_fp8_p32_exact_producer_105fd84009_b2c026f361b4 { + +// 2D TMA descriptor for buffer 'Q_map' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_Q_map(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'Q_map' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'Q_map' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'Q_map' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "Q_map"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'Q_map' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'Q_map' resolved a non-positive global dim"; + TVM_FFI_CHECK(16u <= global_dim[1], ValueError) + << "TMA box (128, 16) exceeds resolved global dims for 'Q_map'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'Q_map' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'Q_map' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'Q_map' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 16u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'Q_map') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'K_map' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_K_map(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'K_map' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'K_map' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'K_map' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "K_map"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'K_map' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'K_map' resolved a non-positive global dim"; + TVM_FFI_CHECK(32u <= global_dim[1], ValueError) + << "TMA box (128, 32) exceeds resolved global dims for 'K_map'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'K_map' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'K_map' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'K_map' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 32u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'K_map') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'V_map' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_V_map(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'V_map' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'V_map' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'V_map' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "V_map"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'V_map' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'V_map' resolved a non-positive global dim"; + TVM_FFI_CHECK(32u <= global_dim[1], ValueError) + << "TMA box (128, 32) exceeds resolved global dims for 'V_map'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'V_map' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'V_map' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'V_map' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 32u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'V_map') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_Q_map, TensorView arg_K_map, TensorView arg_V_map, TensorView arg_page_table, TensorView arg_seq_lens, TensorView arg_partial_O, TensorView arg_partial_stats, TensorView arg_O, TensorView arg_completion, int64_t arg_q_len, int64_t arg_page_table_stride, int64_t arg_max_num_ctas_q, int64_t arg_max_num_ctas_kv, double arg_bmm1_scale_log2, double arg_bmm2_scale, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_Q_map.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_Q_map, "Q_map"); + CheckDtype(arg_Q_map, "Q_map", 1, 8, 1); + CheckCudaTensor(arg_K_map, "K_map"); + CheckDtype(arg_K_map, "K_map", 1, 8, 1); + CheckCudaTensor(arg_V_map, "V_map"); + CheckDtype(arg_V_map, "V_map", 1, 8, 1); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckDtype(arg_seq_lens, "seq_lens", 0, 32, 1); + CheckContiguous(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); + CheckContiguous(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_stats, "partial_stats"); + CheckDtype(arg_partial_stats, "partial_stats", 2, 32, 1); + CheckContiguous(arg_partial_stats, "partial_stats"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_completion, "completion"); + CheckDtype(arg_completion, "completion", 1, 32, 1); + CheckContiguous(arg_completion, "completion"); + TVM_FFI_CHECK(arg_q_len >= -2147483648LL && arg_q_len <= 2147483647LL, ValueError) + << "scalar 'q_len' value " << arg_q_len + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_table_stride >= -2147483648LL && arg_page_table_stride <= 2147483647LL, ValueError) + << "scalar 'page_table_stride' value " << arg_page_table_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_num_ctas_q >= -2147483648LL && arg_max_num_ctas_q <= 2147483647LL, ValueError) + << "scalar 'max_num_ctas_q' value " << arg_max_num_ctas_q + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_num_ctas_kv >= -2147483648LL && arg_max_num_ctas_kv <= 2147483647LL, ValueError) + << "scalar 'max_num_ctas_kv' value " << arg_max_num_ctas_kv + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_K_map, arg_Q_map, "K_map", "Q_map"); + CheckSameCudaDevice(arg_V_map, arg_Q_map, "V_map", "Q_map"); + CheckSameCudaDevice(arg_page_table, arg_Q_map, "page_table", "Q_map"); + CheckSameCudaDevice(arg_seq_lens, arg_Q_map, "seq_lens", "Q_map"); + CheckSameCudaDevice(arg_partial_O, arg_Q_map, "partial_O", "Q_map"); + CheckSameCudaDevice(arg_partial_stats, arg_Q_map, "partial_stats", "Q_map"); + CheckSameCudaDevice(arg_O, arg_Q_map, "O", "Q_map"); + CheckSameCudaDevice(arg_completion, arg_Q_map, "completion", "Q_map"); + CheckCurrentCudaDevice(arg_Q_map, "Q_map"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + CUtensorMap p_Q_map = EncodeTma_Q_map(arg_Q_map); + CUtensorMap p_K_map = EncodeTma_K_map(arg_K_map); + CUtensorMap p_V_map = EncodeTma_V_map(arg_V_map); + void* p_page_table = arg_page_table.data_ptr(); + void* p_seq_lens = arg_seq_lens.data_ptr(); + void* p_partial_O = arg_partial_O.data_ptr(); + void* p_partial_stats = arg_partial_stats.data_ptr(); + void* p_O = arg_O.data_ptr(); + void* p_completion = arg_completion.data_ptr(); + int32_t v_q_len = (int32_t)arg_q_len; + int32_t v_page_table_stride = (int32_t)arg_page_table_stride; + int32_t v_max_num_ctas_q = (int32_t)arg_max_num_ctas_q; + int32_t v_max_num_ctas_kv = (int32_t)arg_max_num_ctas_kv; + float v_bmm1_scale_log2 = (float)arg_bmm1_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + void* kargs[] = {&p_Q_map, &p_K_map, &p_V_map, &p_page_table, &p_seq_lens, &p_partial_O, &p_partial_stats, &p_O, &p_completion, &v_q_len, &v_page_table_stride, &v_max_num_ctas_q, &v_max_num_ctas_kv, &v_bmm1_scale_log2, &v_bmm2_scale}; + + static auto kernel = EmbedCubinModule_fp8_p32_exact_producer_105fd84009::Global()->mod.GetKernel("kernel_fp8_p32_exact_producer"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_Q_map.device().device_id, 178560, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(512u, 1u, 1u); + + if (use_oversized_smem) { + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[1]; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 178560u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[1]; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 178560u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); + + } else { + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.Launch(kargs, grid, block, stream, 178560u)); + } +} +} // namespace variant_fp8_p32_exact_producer_105fd84009_b2c026f361b4 + +int SelectVariant(int32_t batch, int32_t q_len, int32_t page_table_stride, int32_t max_num_ctas_q, int32_t max_num_ctas_kv, float bmm1_scale_log2, float bmm2_scale) { + if (((batch > 0) && (q_len > 0))) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_query, tvm::ffi::TensorView arg_kv, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_completion, tvm::ffi::TensorView arg_ws_partial_output, tvm::ffi::TensorView arg_ws_partial_stats, int32_t batch, int32_t q_len, int32_t page_table_stride, int32_t max_num_ctas_q, int32_t max_num_ctas_kv, float bmm1_scale_log2, float bmm2_scale, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_query, "query"); + CheckCudaTensor(arg_kv, "kv"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_completion, "completion"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(batch, q_len, page_table_stride, max_num_ctas_q, max_num_ctas_kv, bmm1_scale_log2, bmm2_scale); + switch (v) { + case 0: { + CheckCudaTensor(arg_ws_partial_output, "partial_output"); + CheckDtype(arg_ws_partial_output, "partial_output", 4, 16, 1); + TVM_FFI_CHECK(arg_ws_partial_output.numel() == (int64_t)((((batch * 32) * q_len)) * (2) * (16) * (128)), ValueError) + << "workspace 'partial_output' numel " << arg_ws_partial_output.numel() + << " != expected " << (int64_t)((((batch * 32) * q_len)) * (2) * (16) * (128)); + CheckCudaTensor(arg_ws_partial_stats, "partial_stats"); + CheckDtype(arg_ws_partial_stats, "partial_stats", 2, 32, 1); + TVM_FFI_CHECK(arg_ws_partial_stats.numel() == (int64_t)((((batch * 32) * q_len)) * (2) * (16) * (2)), ValueError) + << "workspace 'partial_stats' numel " << arg_ws_partial_stats.numel() + << " != expected " << (int64_t)((((batch * 32) * q_len)) * (2) * (16) * (2)); + { + int64_t gx = (q_len * 2); + int64_t gy = 32; + int64_t gz = batch; + mla_host_shim::variant_fp8_p32_exact_producer_105fd84009_b2c026f361b4::Run(arg_query, arg_kv, arg_kv, arg_page_table, arg_seq_lens, arg_ws_partial_output, arg_ws_partial_stats, arg_output, arg_completion, q_len, page_table_stride, max_num_ctas_q, max_num_ctas_kv, bmm1_scale_log2, bmm2_scale, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_fp8_p32_qk_l2 shape (batch=" << batch << ", q_len=" << q_len << ", page_table_stride=" << page_table_stride << ", max_num_ctas_q=" << max_num_ctas_q << ", max_num_ctas_kv=" << max_num_ctas_kv << ", bmm1_scale_log2=" << bmm1_scale_log2 << ", bmm2_scale=" << bmm2_scale << ")"; + } +} + +int64_t RouteOf(int32_t batch, int32_t q_len, int32_t page_table_stride, int32_t max_num_ctas_q, int32_t max_num_ctas_kv, float bmm1_scale_log2, float bmm2_scale) { return SelectVariant(batch, q_len, page_table_stride, max_num_ctas_q, max_num_ctas_kv, bmm1_scale_log2, bmm2_scale); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_fp8_page64_pdl.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_fp8_page64_pdl.cpp new file mode 100644 index 00000000000..152328b83f6 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_fp8_page64_pdl.cpp @@ -0,0 +1,767 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(page64_keeps_gmem_separate_reduce_cea01a87e9); +TVM_FFI_EMBED_CUBIN(trtllm_mla_decode_fp8_6fe1e0de09); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +struct TmaDeviceArena { + static constexpr size_t kSlotsPerChunk = 256; + static constexpr size_t kMaxSlots = 4096; + std::vector chunks; + size_t used = 0; +}; + +// Immutable, process-lifetime device tensor-map slots for the pointer ABI. +// A slot is never rewritten: different descriptor bytes always get a new +// address, so concurrent streams cannot observe a partially updated map. The +// chunked arena caps storage at 512 KiB per CUDA context in this host module. +static inline void* TmaDeviceSlot( + const CUtensorMap& tm, + int device_id, + cudaStream_t stream) { + static std::mutex mu; + static auto* slots = new std::unordered_map(); + static auto* arenas = new std::unordered_map(); + + // Device allocations are context-owned. Resolve and validate the active + // context before cache lookup so a warm entry can never bypass the same + // checks as a cold entry or leak a pointer across contexts on one device. + CUcontext current_context = nullptr; + CUresult result = cuCtxGetCurrent(¤t_context); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_context != nullptr, RuntimeError) + << "pointer TMA ABI requires an active CUDA context: CUresult=" + << static_cast(result); + CUdevice current_device = -1; + result = cuCtxGetDevice(¤t_device); + TVM_FFI_CHECK(result == CUDA_SUCCESS && current_device == device_id, RuntimeError) + << "TMA descriptor device mismatch: current=" << current_device + << ", tensor=" << device_id; + + std::string key = + std::to_string(reinterpret_cast(current_context)); + key.push_back(':'); + key.append(reinterpret_cast(&tm), sizeof(CUtensorMap)); + std::lock_guard lock(mu); + + // CUcontext is an opaque, recyclable handle. Before trusting a warm arena, + // prove that its first allocation still belongs to this live context. A + // destroyed context makes the pointer query fail; a subsequently reused + // handle must therefore discard the stale pointer and descriptor entries. + auto arena_it = arenas->find(current_context); + if (arena_it != arenas->end() && !arena_it->second.chunks.empty()) { + CUcontext allocation_context = nullptr; + result = cuPointerGetAttribute( + &allocation_context, + CU_POINTER_ATTRIBUTE_CONTEXT, + arena_it->second.chunks.front()); + if (result != CUDA_SUCCESS || allocation_context != current_context) { + const std::string prefix = + std::to_string(reinterpret_cast(current_context)) + ":"; + for (auto slot_it = slots->begin(); slot_it != slots->end();) { + if (slot_it->first.compare(0, prefix.size(), prefix) == 0) { + slot_it = slots->erase(slot_it); + } else { + ++slot_it; + } + } + arenas->erase(arena_it); + } + } + + auto it = slots->find(key); + if (it != slots->end()) return it->second; + + CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; + result = cuStreamIsCapturing( + reinterpret_cast(stream), &capture_status); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuStreamIsCapturing for TMA descriptor slot failed: CUresult=" + << static_cast(result); + TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) + << "pointer TMA ABI cannot create a new device descriptor slot inside " + "CUDA Graph capture; prewarm this exact tensor/layout binding or " + "compile with tma_abi='grid_constant'"; + + TmaDeviceArena& arena = (*arenas)[current_context]; + TVM_FFI_CHECK(arena.used < TmaDeviceArena::kMaxSlots, RuntimeError) + << "pointer TMA ABI exhausted its immutable descriptor arena in CUDA " + "context " << current_context << " on device " << device_id + << " (capacity=" << TmaDeviceArena::kMaxSlots + << "); reuse tensor/layout bindings or compile with tma_abi='grid_constant'"; + if (arena.used % TmaDeviceArena::kSlotsPerChunk == 0) { + CUdeviceptr chunk = 0; + result = cuMemAlloc( + &chunk, + TmaDeviceArena::kSlotsPerChunk * sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemAlloc for TMA descriptor arena failed: CUresult=" + << static_cast(result); + arena.chunks.push_back(chunk); + } + size_t chunk_index = arena.used / TmaDeviceArena::kSlotsPerChunk; + size_t slot_index = arena.used % TmaDeviceArena::kSlotsPerChunk; + CUdeviceptr dev = arena.chunks[chunk_index] + + slot_index * sizeof(CUtensorMap); + result = cuMemcpyHtoD(dev, &tm, sizeof(CUtensorMap)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "cuMemcpyHtoD for TMA descriptor slot failed: CUresult=" + << static_cast(result); + ++arena.used; + void* pointer = reinterpret_cast(static_cast(dev)); + (*slots)[key] = pointer; + return pointer; +} + +namespace variant_trtllm_mla_decode_fp8_6fe1e0de09_c0f86ac47133 { + +// 2D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'tmap_q' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[1], ValueError) + << "TMA box (128, 64) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 64u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'tmap_k' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[1], ValueError) + << "TMA box (128, 64) exceeds resolved global dims for 'tmap_k'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 64u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'tmap_v' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[1], ValueError) + << "TMA box (128, 64) exceeds resolved global dims for 'tmap_v'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 64u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_v') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_k, TensorView arg_tmap_v, TensorView arg_partial_O, TensorView arg_partial_stats, TensorView arg_O, TensorView arg_LSE, TensorView arg_seq_lens_kv, TensorView arg_work_batch_indices, TensorView arg_page_table, double arg_softmax_scale_log2, double arg_bmm2_scale, int64_t arg_num_split, int64_t arg_num_queries, int64_t arg_max_pages_per_seq, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_k, "tmap_k"); + CheckDtype(arg_tmap_k, "tmap_k", 1, 8, 1); + CheckContiguous(arg_tmap_k, "tmap_k"); + CheckCudaTensor(arg_tmap_v, "tmap_v"); + CheckDtype(arg_tmap_v, "tmap_v", 1, 8, 1); + CheckContiguous(arg_tmap_v, "tmap_v"); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); + CheckContiguous(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_stats, "partial_stats"); + CheckDtype(arg_partial_stats, "partial_stats", 2, 32, 1); + CheckContiguous(arg_partial_stats, "partial_stats"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_LSE, "LSE"); + CheckDtype(arg_LSE, "LSE", 2, 32, 1); + CheckContiguous(arg_LSE, "LSE"); + CheckCudaTensor(arg_seq_lens_kv, "seq_lens_kv"); + CheckDtype(arg_seq_lens_kv, "seq_lens_kv", 0, 32, 1); + CheckContiguous(arg_seq_lens_kv, "seq_lens_kv"); + CheckCudaTensor(arg_work_batch_indices, "work_batch_indices"); + CheckDtype(arg_work_batch_indices, "work_batch_indices", 0, 32, 1); + CheckContiguous(arg_work_batch_indices, "work_batch_indices"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) + << "scalar 'num_split' value " << arg_num_split + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_queries >= -2147483648LL && arg_num_queries <= 2147483647LL, ValueError) + << "scalar 'num_queries' value " << arg_num_queries + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_pages_per_seq >= -2147483648LL && arg_max_pages_per_seq <= 2147483647LL, ValueError) + << "scalar 'max_pages_per_seq' value " << arg_max_pages_per_seq + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_k, arg_tmap_q, "tmap_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_v, arg_tmap_q, "tmap_v", "tmap_q"); + CheckSameCudaDevice(arg_partial_O, arg_tmap_q, "partial_O", "tmap_q"); + CheckSameCudaDevice(arg_partial_stats, arg_tmap_q, "partial_stats", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_LSE, arg_tmap_q, "LSE", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens_kv, arg_tmap_q, "seq_lens_kv", "tmap_q"); + CheckSameCudaDevice(arg_work_batch_indices, arg_tmap_q, "work_batch_indices", "tmap_q"); + CheckSameCudaDevice(arg_page_table, arg_tmap_q, "page_table", "tmap_q"); + CheckCurrentCudaDevice(arg_tmap_q, "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + void* p_tmap_q = TmaDeviceSlot(h_tmap_q, arg_tmap_q.device().device_id, stream); + CUtensorMap h_tmap_k = EncodeTma_tmap_k(arg_tmap_k); + void* p_tmap_k = TmaDeviceSlot(h_tmap_k, arg_tmap_k.device().device_id, stream); + CUtensorMap h_tmap_v = EncodeTma_tmap_v(arg_tmap_v); + void* p_tmap_v = TmaDeviceSlot(h_tmap_v, arg_tmap_v.device().device_id, stream); + void* p_partial_O = arg_partial_O.data_ptr(); + void* p_partial_stats = arg_partial_stats.data_ptr(); + void* p_O = arg_O.data_ptr(); + void* p_LSE = arg_LSE.data_ptr(); + void* p_seq_lens_kv = arg_seq_lens_kv.data_ptr(); + void* p_work_batch_indices = arg_work_batch_indices.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + float v_softmax_scale_log2 = (float)arg_softmax_scale_log2; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_num_split = (int32_t)arg_num_split; + int32_t v_num_queries = (int32_t)arg_num_queries; + int32_t v_max_pages_per_seq = (int32_t)arg_max_pages_per_seq; + void* kargs[] = {&p_tmap_q, &p_tmap_k, &p_tmap_v, &p_partial_O, &p_partial_stats, &p_O, &p_LSE, &p_seq_lens_kv, &p_work_batch_indices, &p_page_table, &v_softmax_scale_log2, &v_bmm2_scale, &v_num_split, &v_num_queries, &v_max_pages_per_seq}; + + static auto kernel = EmbedCubinModule_trtllm_mla_decode_fp8_6fe1e0de09::Global()->mod.GetKernel("kernel_trtllm_mla_decode_fp8"); + static signed char mla_smem_mode_cache[64] = {0}; + const bool use_oversized_smem = MlaConfigureDynamicSmem( + kernel, (int)arg_tmap_q.device().device_id, 190976, + mla_smem_mode_cache, 64); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(384u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[3]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[n].value.clusterDim.x = 2u; + attrs[n].value.clusterDim.y = 1u; + attrs[n].value.clusterDim.z = 1u; + ++n; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE; + attrs[n].value.sharedMemoryMode = CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY; + ++n; + } +#endif + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 190976u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[3]; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; +#if MLA_HAS_OVERSIZED_SMEM + if (use_oversized_smem) { + attrs[n].id = cudaLaunchAttributeSharedMemoryMode; + attrs[n].val.sharedMemoryMode = cudaSharedMemoryModeAllowOversizedSharedMemory; + ++n; + } +#endif + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 190976u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_trtllm_mla_decode_fp8_6fe1e0de09_c0f86ac47133 + +namespace variant_page64_keeps_gmem_separate_reduce_cea01a87e9_dde077f61656 { + +void Run(TensorView arg_partial_O, TensorView arg_partial_stats, TensorView arg_O, TensorView arg_LSE, int64_t arg_num_split, int64_t arg_num_reduce_ctas, int64_t arg_num_queries, double arg_bmm2_scale, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_partial_O.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); + CheckContiguous(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_stats, "partial_stats"); + CheckDtype(arg_partial_stats, "partial_stats", 2, 32, 1); + CheckContiguous(arg_partial_stats, "partial_stats"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_LSE, "LSE"); + CheckDtype(arg_LSE, "LSE", 2, 32, 1); + CheckContiguous(arg_LSE, "LSE"); + TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) + << "scalar 'num_split' value " << arg_num_split + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_reduce_ctas >= -2147483648LL && arg_num_reduce_ctas <= 2147483647LL, ValueError) + << "scalar 'num_reduce_ctas' value " << arg_num_reduce_ctas + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_queries >= -2147483648LL && arg_num_queries <= 2147483647LL, ValueError) + << "scalar 'num_queries' value " << arg_num_queries + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_partial_stats, arg_partial_O, "partial_stats", "partial_O"); + CheckSameCudaDevice(arg_O, arg_partial_O, "O", "partial_O"); + CheckSameCudaDevice(arg_LSE, arg_partial_O, "LSE", "partial_O"); + CheckCurrentCudaDevice(arg_partial_O, "partial_O"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + void* p_partial_O = arg_partial_O.data_ptr(); + void* p_partial_stats = arg_partial_stats.data_ptr(); + void* p_O = arg_O.data_ptr(); + void* p_LSE = arg_LSE.data_ptr(); + int32_t v_num_split = (int32_t)arg_num_split; + int32_t v_num_reduce_ctas = (int32_t)arg_num_reduce_ctas; + int32_t v_num_queries = (int32_t)arg_num_queries; + float v_bmm2_scale = (float)arg_bmm2_scale; + void* kargs[] = {&p_partial_O, &p_partial_stats, &p_O, &p_LSE, &v_num_split, &v_num_reduce_ctas, &v_num_queries, &v_bmm2_scale}; + + static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(page64_keeps_gmem_separate_reduce_cea01a87e9, "kernel_page64_keeps_gmem_separate_reduce"); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(512u, 1u, 1u); + + // Extended launch — mirrors CUDAKernel.prepare_launch_cluster attributes. + tvm::ffi::cuda_api::LaunchConfig config; + int n = 0; +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + CUlaunchAttribute attrs[1]; + attrs[n].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + attrs[n].value.programmaticStreamSerializationAllowed = 1; + ++n; + config.gridDimX = grid.x; + config.gridDimY = grid.y; + config.gridDimZ = grid.z; + config.blockDimX = block.x; + config.blockDimY = block.y; + config.blockDimZ = block.z; + config.sharedMemBytes = 0u; + config.hStream = stream; + config.attrs = attrs; + config.numAttrs = n; +#else + cudaLaunchAttribute attrs[1]; + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; + config.gridDim = {grid.x, grid.y, grid.z}; + config.blockDim = {block.x, block.y, block.z}; + config.dynamicSmemBytes = 0u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; +#endif + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.LaunchEx(kargs, config)); +} +} // namespace variant_page64_keeps_gmem_separate_reduce_cea01a87e9_dde077f61656 + +int SelectVariant(float softmax_scale_log2, float bmm2_scale, int32_t num_split, int32_t num_queries, int32_t max_pages_per_seq, int32_t num_reduce_ctas) { + if ((num_split > 1)) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_query, tvm::ffi::TensorView arg_kv, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_lse, tvm::ffi::TensorView arg_seq_lens, tvm::ffi::TensorView arg_work_batch_indices, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_ws_partial_output, tvm::ffi::TensorView arg_ws_partial_stats, float softmax_scale_log2, float bmm2_scale, int32_t num_split, int32_t num_queries, int32_t max_pages_per_seq, int32_t num_reduce_ctas, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_query, "query"); + CheckCudaTensor(arg_kv, "kv"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_lse, "lse"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_work_batch_indices, "work_batch_indices"); + CheckCudaTensor(arg_page_table, "page_table"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(softmax_scale_log2, bmm2_scale, num_split, num_queries, max_pages_per_seq, num_reduce_ctas); + switch (v) { + case 0: { + CheckCudaTensor(arg_ws_partial_output, "partial_output"); + CheckDtype(arg_ws_partial_output, "partial_output", 4, 16, 1); + TVM_FFI_CHECK(arg_ws_partial_output.numel() == (int64_t)((num_queries) * (num_split) * (128) * (512)), ValueError) + << "workspace 'partial_output' numel " << arg_ws_partial_output.numel() + << " != expected " << (int64_t)((num_queries) * (num_split) * (128) * (512)); + CheckCudaTensor(arg_ws_partial_stats, "partial_stats"); + CheckDtype(arg_ws_partial_stats, "partial_stats", 2, 32, 1); + TVM_FFI_CHECK(arg_ws_partial_stats.numel() == (int64_t)((num_queries) * (num_split) * (128) * (2)), ValueError) + << "workspace 'partial_stats' numel " << arg_ws_partial_stats.numel() + << " != expected " << (int64_t)((num_queries) * (num_split) * (128) * (2)); + { + int64_t gx = (num_split * 2); + int64_t gy = num_queries; + int64_t gz = 1; + mla_host_shim::variant_trtllm_mla_decode_fp8_6fe1e0de09_c0f86ac47133::Run(arg_query, arg_kv, arg_kv, arg_ws_partial_output, arg_ws_partial_stats, arg_output, arg_lse, arg_seq_lens, arg_work_batch_indices, arg_page_table, softmax_scale_log2, bmm2_scale, num_split, num_queries, max_pages_per_seq, gx, gy, gz, stream); + } + { + int64_t gx = num_reduce_ctas; + int64_t gy = 2; + int64_t gz = num_queries; + mla_host_shim::variant_page64_keeps_gmem_separate_reduce_cea01a87e9_dde077f61656::Run(arg_ws_partial_output, arg_ws_partial_stats, arg_output, arg_lse, num_split, num_reduce_ctas, num_queries, bmm2_scale, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_fp8_page64_pdl shape (softmax_scale_log2=" << softmax_scale_log2 << ", bmm2_scale=" << bmm2_scale << ", num_split=" << num_split << ", num_queries=" << num_queries << ", max_pages_per_seq=" << max_pages_per_seq << ", num_reduce_ctas=" << num_reduce_ctas << ")"; + } +} + +int64_t RouteOf(float softmax_scale_log2, float bmm2_scale, int32_t num_split, int32_t num_queries, int32_t max_pages_per_seq, int32_t num_reduce_ctas) { return SelectVariant(softmax_scale_log2, bmm2_scale, num_split, num_queries, max_pages_per_seq, num_reduce_ctas); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_fp8_tail.cpp b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_fp8_tail.cpp new file mode 100644 index 00000000000..f551a383035 --- /dev/null +++ b/csrc/cake_trtllm_mla_blackwell/generated/host/sm_103a_152/cake_mla_fp8_tail.cpp @@ -0,0 +1,409 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Generated source; do not edit manually. +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TVM_FFI_EMBED_CUBIN(full_abi_tail_fp8_26b36022eb); + +namespace mla_host_shim { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t MlaDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +#include + +// CUDA 13.4 adds an oversized shared-memory mode (function/launch attribute +// SHARED_MEMORY_MODE = 3) that lets a kernel exceed the standard +// MaxSharedMemoryPerBlockOptin ceiling up to device attribute 150 +// (OversizedSharedMemoryPerBlock), at the cost of an 8 KiB L1 carveout. +#if (TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDA_VERSION) && CUDA_VERSION >= 13040) || \ + (!TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API && defined(CUDART_VERSION) && CUDART_VERSION >= 13040) +#define MLA_HAS_OVERSIZED_SMEM 1 +#else +#define MLA_HAS_OVERSIZED_SMEM 0 +#endif + +// Per-device dynamic-SMEM opt-in. Returns true when launches on this device +// must carry the ALLOW_OVERSIZED shared-memory-mode launch attribute: the +// request exceeds the device's standard opt-in ceiling +// (MaxSharedMemoryPerBlockOptin, attribute 97) but fits the oversized ceiling +// (OversizedSharedMemoryPerBlock, attribute 150). Within the standard ceiling +// it sets MAX_DYNAMIC_SHARED_SIZE_BYTES for the device, mirroring +// CubinKernel::SetMaxDynamicSharedMemory. cache: 0 = unresolved, +// 1 = standard opt-in done, 2 = oversized mode required. +inline bool MlaConfigureDynamicSmem(tvm::ffi::CubinKernel& kernel, int device_id, + int smem_bytes, signed char* cache, int cache_len) { + namespace cuda_api = tvm::ffi::cuda_api; + TVM_FFI_CHECK(device_id >= 0 && device_id < cache_len, RuntimeError) + << "dynamic-SMEM opt-in cache does not cover cuda:" << device_id; + if (cache[device_id] != 0) { + return cache[device_id] == 2; + } + auto device = cuda_api::GetDeviceHandle(device_id); + int optin_max = 0; + cuda_api::ResultType err = cuda_api::GetDeviceAttribute( + &optin_max, + /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN / + cudaDevAttrMaxSharedMemoryPerBlockOptin */ + cuda_api::DeviceAttrType(97), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "querying MaxSharedMemoryPerBlockOptin failed for cuda:" << device_id; + if (smem_bytes <= optin_max) { + err = cuda_api::SetKernelMaxDynamicSharedMem(kernel.GetHandle(), smem_bytes, device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "MAX_DYNAMIC_SHARED_SIZE_BYTES=" << smem_bytes << " rejected for cuda:" << device_id; + cache[device_id] = 1; + return false; + } +#if MLA_HAS_OVERSIZED_SMEM + int oversized_max = 0; + err = cuda_api::GetDeviceAttribute( + &oversized_max, + /* CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK / + cudaDevAttrOversizedSharedMemoryPerBlock */ + cuda_api::DeviceAttrType(150), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id << " and the oversized ceiling is " + << oversized_max << " B"; + // Also set the function-level shared-memory mode: profiler-instrumented + // launches (ncu) honor only the function-level attribute for oversized + // dynamic SMEM and fail with LaunchFailed on the launch attribute alone. +#if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API + err = cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, + CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY, + kernel.GetHandle(), device); + TVM_FFI_CHECK(err == cuda_api::kSuccess, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (error " << static_cast(err) << ")"; +#else + // The process cudart may predate 13.4 (e.g. torch's pip cudart 13.0) and + // reject cudaFuncAttributeSharedMemoryMode with cudaErrorInvalidValue even + // though the driver supports the mode, so call the driver entry point + // directly: cudaKernel_t is interchangeable with CUkernel, the runtime + // device ordinal is the CUdevice, and libcuda.so.1 is already loaded. + { + using MlaCuKernelSetAttributeFn = int (*)(int attrib, int val, void* kernel, int dev); + static const auto mla_cu_kernel_set_attribute = + reinterpret_cast(dlsym(RTLD_DEFAULT, "cuKernelSetAttribute")); + TVM_FFI_CHECK(mla_cu_kernel_set_attribute != nullptr, RuntimeError) + << "cuKernelSetAttribute not resolvable while enabling the oversized shared-memory mode"; + int drv_err = mla_cu_kernel_set_attribute( + /* CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE */ 17, + /* CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY */ 3, + reinterpret_cast(kernel.GetHandle()), device_id); + TVM_FFI_CHECK(drv_err == 0, RuntimeError) + << "setting SHARED_MEMORY_MODE=ALLOW_OVERSIZED failed for cuda:" << device_id + << " (driver error " << drv_err << ")"; + } +#endif + cache[device_id] = 2; + return true; +#else + TVM_FFI_THROW(RuntimeError) + << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling (" + << optin_max << " B) on cuda:" << device_id + << " and this CUDA toolkit predates the 13.4 oversized shared-memory mode"; +#endif +} + +namespace variant_full_abi_tail_fp8_26b36022eb_f0c2469a092e { + +void Run(TensorView arg_Q, TensorView arg_KV, TensorView arg_fp8_lut, TensorView arg_page_table, TensorView arg_sparse_indices, TensorView arg_row_batches, TensorView arg_row_seq_lens, TensorView arg_O, TensorView arg_LSE, TensorView arg_sinks, int64_t arg_num_heads, int64_t arg_qk_dim, int64_t arg_value_dim, int64_t arg_kv_stride, int64_t arg_page_size, int64_t arg_page_table_width, int64_t arg_sparse_width, int64_t arg_use_sparse, double arg_softmax_scale, double arg_bmm2_scale, int64_t arg_enable_sink, int64_t arg_write_lse, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_Q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_Q, "Q"); + CheckDtype(arg_Q, "Q", 1, 8, 1); + CheckContiguous(arg_Q, "Q"); + CheckCudaTensor(arg_KV, "KV"); + CheckDtype(arg_KV, "KV", 1, 8, 1); + CheckContiguous(arg_KV, "KV"); + CheckCudaTensor(arg_fp8_lut, "fp8_lut"); + CheckDtype(arg_fp8_lut, "fp8_lut", 2, 32, 1); + CheckContiguous(arg_fp8_lut, "fp8_lut"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckDtype(arg_page_table, "page_table", 0, 32, 1); + CheckContiguous(arg_page_table, "page_table"); + CheckCudaTensor(arg_sparse_indices, "sparse_indices"); + CheckDtype(arg_sparse_indices, "sparse_indices", 0, 32, 1); + CheckContiguous(arg_sparse_indices, "sparse_indices"); + CheckCudaTensor(arg_row_batches, "row_batches"); + CheckDtype(arg_row_batches, "row_batches", 0, 32, 1); + CheckContiguous(arg_row_batches, "row_batches"); + CheckCudaTensor(arg_row_seq_lens, "row_seq_lens"); + CheckDtype(arg_row_seq_lens, "row_seq_lens", 0, 32, 1); + CheckContiguous(arg_row_seq_lens, "row_seq_lens"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_LSE, "LSE"); + CheckDtype(arg_LSE, "LSE", 2, 32, 1); + CheckContiguous(arg_LSE, "LSE"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_qk_dim >= -2147483648LL && arg_qk_dim <= 2147483647LL, ValueError) + << "scalar 'qk_dim' value " << arg_qk_dim + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_value_dim >= -2147483648LL && arg_value_dim <= 2147483647LL, ValueError) + << "scalar 'value_dim' value " << arg_value_dim + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_kv_stride >= -2147483648LL && arg_kv_stride <= 2147483647LL, ValueError) + << "scalar 'kv_stride' value " << arg_kv_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_size >= -2147483648LL && arg_page_size <= 2147483647LL, ValueError) + << "scalar 'page_size' value " << arg_page_size + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_table_width >= -2147483648LL && arg_page_table_width <= 2147483647LL, ValueError) + << "scalar 'page_table_width' value " << arg_page_table_width + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_width >= -2147483648LL && arg_sparse_width <= 2147483647LL, ValueError) + << "scalar 'sparse_width' value " << arg_sparse_width + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_sparse >= -2147483648LL && arg_use_sparse <= 2147483647LL, ValueError) + << "scalar 'use_sparse' value " << arg_use_sparse + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_enable_sink >= -2147483648LL && arg_enable_sink <= 2147483647LL, ValueError) + << "scalar 'enable_sink' value " << arg_enable_sink + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_write_lse >= -2147483648LL && arg_write_lse <= 2147483647LL, ValueError) + << "scalar 'write_lse' value " << arg_write_lse + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_KV, arg_Q, "KV", "Q"); + CheckSameCudaDevice(arg_fp8_lut, arg_Q, "fp8_lut", "Q"); + CheckSameCudaDevice(arg_page_table, arg_Q, "page_table", "Q"); + CheckSameCudaDevice(arg_sparse_indices, arg_Q, "sparse_indices", "Q"); + CheckSameCudaDevice(arg_row_batches, arg_Q, "row_batches", "Q"); + CheckSameCudaDevice(arg_row_seq_lens, arg_Q, "row_seq_lens", "Q"); + CheckSameCudaDevice(arg_O, arg_Q, "O", "Q"); + CheckSameCudaDevice(arg_LSE, arg_Q, "LSE", "Q"); + CheckSameCudaDevice(arg_sinks, arg_Q, "sinks", "Q"); + CheckCurrentCudaDevice(arg_Q, "Q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + void* p_Q = arg_Q.data_ptr(); + void* p_KV = arg_KV.data_ptr(); + void* p_fp8_lut = arg_fp8_lut.data_ptr(); + void* p_page_table = arg_page_table.data_ptr(); + void* p_sparse_indices = arg_sparse_indices.data_ptr(); + void* p_row_batches = arg_row_batches.data_ptr(); + void* p_row_seq_lens = arg_row_seq_lens.data_ptr(); + void* p_O = arg_O.data_ptr(); + void* p_LSE = arg_LSE.data_ptr(); + void* p_sinks = arg_sinks.data_ptr(); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_qk_dim = (int32_t)arg_qk_dim; + int32_t v_value_dim = (int32_t)arg_value_dim; + int32_t v_kv_stride = (int32_t)arg_kv_stride; + int32_t v_page_size = (int32_t)arg_page_size; + int32_t v_page_table_width = (int32_t)arg_page_table_width; + int32_t v_sparse_width = (int32_t)arg_sparse_width; + int32_t v_use_sparse = (int32_t)arg_use_sparse; + float v_softmax_scale = (float)arg_softmax_scale; + float v_bmm2_scale = (float)arg_bmm2_scale; + int32_t v_enable_sink = (int32_t)arg_enable_sink; + int32_t v_write_lse = (int32_t)arg_write_lse; + void* kargs[] = {&p_Q, &p_KV, &p_fp8_lut, &p_page_table, &p_sparse_indices, &p_row_batches, &p_row_seq_lens, &p_O, &p_LSE, &p_sinks, &v_num_heads, &v_qk_dim, &v_value_dim, &v_kv_stride, &v_page_size, &v_page_table_width, &v_sparse_width, &v_use_sparse, &v_softmax_scale, &v_bmm2_scale, &v_enable_sink, &v_write_lse}; + + static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(full_abi_tail_fp8_26b36022eb, "kernel_full_abi_tail_fp8"); + tvm::ffi::dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + tvm::ffi::dim3 block(32u, 1u, 1u); + + TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.Launch(kargs, grid, block, stream, 16384u)); +} +} // namespace variant_full_abi_tail_fp8_26b36022eb_f0c2469a092e + +int SelectVariant(int32_t num_heads, int32_t qk_dim, int32_t value_dim, int32_t kv_stride, int32_t page_size, int32_t page_table_width, int32_t sparse_width, int32_t use_sparse, float softmax_scale, float bmm2_scale, int32_t enable_sink, int32_t write_lse, int32_t num_rows) { + if ((num_rows > 0)) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_query, tvm::ffi::TensorView arg_kv, tvm::ffi::TensorView arg_fp8_lut, tvm::ffi::TensorView arg_page_table, tvm::ffi::TensorView arg_sparse_indices, tvm::ffi::TensorView arg_row_batches, tvm::ffi::TensorView arg_row_seq_lens, tvm::ffi::TensorView arg_output, tvm::ffi::TensorView arg_lse, tvm::ffi::TensorView arg_sinks, int32_t num_heads, int32_t qk_dim, int32_t value_dim, int32_t kv_stride, int32_t page_size, int32_t page_table_width, int32_t sparse_width, int32_t use_sparse, float softmax_scale, float bmm2_scale, int32_t enable_sink, int32_t write_lse, int32_t num_rows, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_query, "query"); + CheckCudaTensor(arg_kv, "kv"); + CheckCudaTensor(arg_fp8_lut, "fp8_lut"); + CheckCudaTensor(arg_page_table, "page_table"); + CheckCudaTensor(arg_sparse_indices, "sparse_indices"); + CheckCudaTensor(arg_row_batches, "row_batches"); + CheckCudaTensor(arg_row_seq_lens, "row_seq_lens"); + CheckCudaTensor(arg_output, "output"); + CheckCudaTensor(arg_lse, "lse"); + CheckCudaTensor(arg_sinks, "sinks"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(num_heads, qk_dim, value_dim, kv_stride, page_size, page_table_width, sparse_width, use_sparse, softmax_scale, bmm2_scale, enable_sink, write_lse, num_rows); + switch (v) { + case 0: { + { + int64_t gx = num_rows; + int64_t gy = num_heads; + int64_t gz = 1; + mla_host_shim::variant_full_abi_tail_fp8_26b36022eb_f0c2469a092e::Run(arg_query, arg_kv, arg_fp8_lut, arg_page_table, arg_sparse_indices, arg_row_batches, arg_row_seq_lens, arg_output, arg_lse, arg_sinks, num_heads, qk_dim, value_dim, kv_stride, page_size, page_table_width, sparse_width, use_sparse, softmax_scale, bmm2_scale, enable_sink, write_lse, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for mla_fp8_tail shape (num_heads=" << num_heads << ", qk_dim=" << qk_dim << ", value_dim=" << value_dim << ", kv_stride=" << kv_stride << ", page_size=" << page_size << ", page_table_width=" << page_table_width << ", sparse_width=" << sparse_width << ", use_sparse=" << use_sparse << ", softmax_scale=" << softmax_scale << ", bmm2_scale=" << bmm2_scale << ", enable_sink=" << enable_sink << ", write_lse=" << write_lse << ", num_rows=" << num_rows << ")"; + } +} + +int64_t RouteOf(int32_t num_heads, int32_t qk_dim, int32_t value_dim, int32_t kv_stride, int32_t page_size, int32_t page_table_width, int32_t sparse_width, int32_t use_sparse, float softmax_scale, float bmm2_scale, int32_t enable_sink, int32_t write_lse, int32_t num_rows) { return SelectVariant(num_heads, qk_dim, value_dim, kv_stride, page_size, page_table_width, sparse_width, use_sparse, softmax_scale, bmm2_scale, enable_sink, write_lse, num_rows); } + +} // namespace mla_host_shim + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, mla_host_shim::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, mla_host_shim::RouteOf); diff --git a/flashinfer/mla/_core.py b/flashinfer/mla/_core.py index 3d146ce4a9c..78490cded6e 100644 --- a/flashinfer/mla/_core.py +++ b/flashinfer/mla/_core.py @@ -3363,7 +3363,7 @@ def _trtllm_batch_decode_with_kv_cache_mla_impl( use_fp16_softmax: Optional[bool] = None, return_lse_base: Optional[Literal["basee", "base2"]] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - r"""Decode MLA with TRTLLM-GEN, CuteDSL, XQA, or SM120/SM121 sparse kernels. + r"""Decode MLA with Blackwell, TRTLLM-GEN, CuteDSL, XQA, or sparse kernels. With ``backend="auto"``, SM100/SM103 devices use TRTLLM-GEN for sparse MLA when ``sparse_mla_top_k > 0``. SM120/SM121 devices use the packed sparse @@ -3433,14 +3433,14 @@ def _trtllm_batch_decode_with_kv_cache_mla_impl( Output tensor. If not provided, it is allocated internally. bmm1_scale : Union[float, torch.Tensor] Fused scale for MLA BMM1. TRTLLM-GEN accepts a FP32 tensor or float. - CuteDSL, XQA, and SM120/SM121 sparse v32/GLM require a float. + Blackwell, CuteDSL, XQA, and SM120/SM121 sparse v32/GLM require a float. bmm2_scale : Union[float, torch.Tensor] Fused scale for MLA BMM2. TRTLLM-GEN accepts a FP32 tensor or float. - CuteDSL and XQA require a float. SM120/SM121 sparse v32/GLM requires - ``1.0``. + Blackwell, CuteDSL, and XQA require a float. SM120/SM121 sparse v32/GLM + requires ``1.0``. sinks : Optional[List[torch.Tensor]] Additional value per head in the denominator of the softmax. - Supported by ``trtllm-gen``, ``cute-dsl``, and ``sparse``. + Supported by ``cake``, ``trtllm-gen``, ``cute-dsl``, and ``sparse``. On ``cute-dsl`` this requires the modular implementation; ``cute_dsl_impl="auto"`` (the default) promotes to modular automatically, and ``cute_dsl_impl="monolithic"`` with sinks set raises @@ -3456,7 +3456,12 @@ def _trtllm_batch_decode_with_kv_cache_mla_impl( backends; ignored by ``cute-dsl``. backend : str = "auto" Implementation backend. Valid values are ``"auto"``, ``"xqa"``, - ``"trtllm-gen"``, ``"cute-dsl"``, and ``"sparse"``. ``"auto"`` + ``"cake"``, ``"trtllm-gen"``, ``"cute-dsl"``, and + ``"sparse"``. + ``"cake"`` explicitly selects the source-level Blackwell semantic + dispatcher. It covers its qualified BF16/FP8 dense, sparse, compact-Q, + sink, LSE, and caller-owned-output envelope and is never selected + automatically. ``"auto"`` chooses ``"trtllm-gen"`` for SM100/SM103 sparse MLA and chooses ``"sparse"`` for SM120/SM121 when ``sparse_mla_top_k > 0``; otherwise SM120/SM121 dense decode uses ``"xqa"``. @@ -3485,8 +3490,8 @@ def _trtllm_batch_decode_with_kv_cache_mla_impl( passing ``True`` to other backends raises ``ValueError``. lse : Optional[torch.Tensor] = None Optional pre-allocated buffer for Log-Sum-Exp values. Supported by - ``trtllm-gen``, ``cute-dsl``, and ``sparse`` backends. Must have - dtype ``torch.float32``. Accepted shapes: + ``cake``, ``trtllm-gen``, ``cute-dsl``, and ``sparse`` backends. Must + have dtype ``torch.float32``. Accepted shapes: * ``[batch_size * q_len_per_request, num_qo_heads]`` (TRTLLM-GEN native; accepted by sparse), or @@ -3498,7 +3503,7 @@ def _trtllm_batch_decode_with_kv_cache_mla_impl( If ``return_lse`` is True and this is None, a buffer will be allocated by the backend. return_lse : bool = False - Whether to return LSE values. Supported by ``trtllm-gen``, + Whether to return LSE values. Supported by ``cake``, ``trtllm-gen``, ``cute-dsl``, and ``sparse`` backends. When True, the function returns ``(out, lse)``. With compact variable Q, LSE is currently supported only by monolithic CuTeDSL. @@ -3535,9 +3540,9 @@ def _trtllm_batch_decode_with_kv_cache_mla_impl( shape ``[batch_size + 1]``, dtype ``torch.int32``. Must be a 1D tensor with at least two entries. When ``max_q_len`` is not provided, this function validates that it starts with 0, ends at ``query.size(0)``, - and is monotonically non-decreasing. Supported by TRTLLM-GEN and the - monolithic CuTeDSL implementation. When provided, ``query`` must have - shape ``[total_q, num_heads, head_dim_qk]``. + and is monotonically non-decreasing. Supported by Blackwell, TRTLLM-GEN, + and the monolithic CuTeDSL implementation. When provided, ``query`` + must have shape ``[total_q, num_heads, head_dim_qk]``. For best performance, provide ``max_q_len`` together with ``cum_seq_lens_q`` to avoid host-side metadata validation. max_q_len : Optional[int] = None @@ -3682,6 +3687,36 @@ def _trtllm_batch_decode_with_kv_cache_mla_impl( use_fp16_softmax, "use_fp16_softmax", query.device ) + if backend == "cake": + from .cake_trtllm_mla_blackwell import trtllm_mla_blackwell_decode + + return trtllm_mla_blackwell_decode( + query=query, + kv_cache=kv_cache, + block_tables=block_tables, + seq_lens=seq_lens, + max_seq_len=max_seq_len, + qk_nope_head_dim=qk_nope_head_dim, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + sparse_mla_top_k=sparse_mla_top_k, + out=out, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + sinks=sinks, + skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale_factor, + enable_pdl=enable_pdl, + uses_shared_paged_kv_idx=uses_shared_paged_kv_idx, + lse=lse, + return_lse=return_lse, + cum_seq_lens_q=cum_seq_lens_q, + max_q_len=max_q_len, + multi_ctas_kv_counter_buffer=multi_ctas_kv_counter_buffer, + sparse_mla_top_k_lens=sparse_mla_top_k_lens, + enable_dcp=enable_dcp, + backend="cake", + ) + if backend == "auto": cc = get_compute_capability(query.device) if cc[0] == 12 and sparse_mla_top_k > 0: diff --git a/flashinfer/mla/cake_trtllm_mla_blackwell.py b/flashinfer/mla/cake_trtllm_mla_blackwell.py new file mode 100644 index 00000000000..c78eca95009 --- /dev/null +++ b/flashinfer/mla/cake_trtllm_mla_blackwell.py @@ -0,0 +1,1494 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import functools +import hashlib +import json +import math +import os +import re +import shutil +import subprocess +import threading +from collections import OrderedDict +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import Any, Callable, Literal, Optional, Union + +import torch +from filelock import FileLock +from tvm_ffi import cpp + +from ..jit import env as jit_env +from ..jit.core import logger +from ..jit.cpp_ext import is_cuda_version_at_least +from ..utils import get_compute_capability, get_device_sm_count, log2e + +# --- Source-built domain loader (formerly flashinfer/jit/cake_trtllm_mla_blackwell.py) --- + +_TARGETS = { + "sm_100a_148": ("sm_100a", 148), + "sm_100a_152": ("sm_100a", 152), + "sm_103a_148": ("sm_103a", 148), + "sm_103a_152": ("sm_103a", 152), +} +_TARGET_ORDER = tuple(_TARGETS) +_ARCH_CAPABILITIES = {"sm_100a": (10, 0), "sm_103a": (10, 3)} +_SOURCE_CATALOG_RELATIVE_PATH = Path("generated") / "cake_source_catalog.json" +_DOMAIN_DEVICE_COUNTS = { + "mla_bf16_vquarter": 1, + "mla_bf16_vhalf": 1, + "mla_bf16_unsplit": 1, + "mla_bf16_clc": 8, + "mla_bf16_tail": 1, + "mla_fp8_tail": 1, + "mla_fp8_p32_qk_l2": 1, + "mla_fp8_page64_pdl": 2, + "mla_bf16_native_split8_pdl": 2, +} +_DOMAIN_ORDER = tuple(_DOMAIN_DEVICE_COUNTS) +_EXPORTED_COMPILE_FLAGS = ["--use_fast_math"] +_HEX_SHA256 = re.compile(r"[0-9a-f]{64}\Z") +_C_IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") + + +def _source_dir() -> Path: + """Locate the generated MLA source package in installs and checkouts.""" + + packaged = jit_env.FLASHINFER_CSRC_DIR / "cake_trtllm_mla_blackwell" + if packaged.is_dir(): + return packaged + return Path(__file__).resolve().parents[2] / "csrc" / "cake_trtllm_mla_blackwell" + + +def _source_record( + value: object, + *, + domain: str, + kind: str, + target: str | None = None, + index: int | None = None, +) -> Mapping[str, object]: + location_parts = [domain] + if target is not None: + location_parts.append(target) + location_parts.append(kind if index is None else f"{kind}[{index}]") + location = "/".join(location_parts) + expected_keys = {"path", "sha256"} + if kind == "device_source": + expected_keys.update({"module_ident", "compile_flags"}) + if not isinstance(value, dict) or set(value) != expected_keys: + raise RuntimeError( + f"TRT-LLM MLA Blackwell catalog {location} schema is invalid" + ) + + relative = value["path"] + sha256 = value["sha256"] + suffix = ".cpp" if kind == "host_source" else ".cu" + if not isinstance(relative, str): + raise RuntimeError(f"TRT-LLM MLA Blackwell catalog {location} path is invalid") + path = PurePosixPath(relative) + if ( + not relative + or "\\" in relative + or path.is_absolute() + or ".." in path.parts + or path.as_posix() != relative + or path.suffix != suffix + ): + raise RuntimeError( + f"TRT-LLM MLA Blackwell catalog {location} path is noncanonical: {relative!r}" + ) + if not isinstance(sha256, str) or _HEX_SHA256.fullmatch(sha256) is None: + raise RuntimeError( + f"TRT-LLM MLA Blackwell catalog {location} sha256 is invalid" + ) + + if kind == "device_source": + module_ident = value["module_ident"] + compile_flags = value["compile_flags"] + if ( + not isinstance(module_ident, str) + or _C_IDENTIFIER.fullmatch(module_ident) is None + ): + raise RuntimeError( + f"TRT-LLM MLA Blackwell catalog {location} module identity is invalid" + ) + if compile_flags != _EXPORTED_COMPILE_FLAGS: + raise RuntimeError( + f"TRT-LLM MLA Blackwell catalog {location} compile flags differ from " + f"the exported contract: {compile_flags!r}" + ) + return value + + +@functools.cache +def _source_catalog() -> Mapping[str, object]: + """Load and validate the exact physical-target generated-source catalog.""" + + catalog_path = _source_dir() / _SOURCE_CATALOG_RELATIVE_PATH + if not catalog_path.is_file(): + raise RuntimeError( + f"TRT-LLM MLA Blackwell generated-source catalog is missing: {catalog_path}" + ) + try: + catalog = json.loads(catalog_path.read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError) as error: + raise RuntimeError( + f"TRT-LLM MLA Blackwell generated-source catalog is unreadable: {catalog_path}" + ) from error + if not isinstance(catalog, dict) or set(catalog) != { + "schema_version", + "target_order", + "targets", + "domain_order", + "domains", + }: + raise RuntimeError( + "TRT-LLM MLA Blackwell generated-source catalog schema is invalid" + ) + target_order = catalog["target_order"] + canonical_targets = [ + target + for target in _TARGET_ORDER + if isinstance(target_order, list) and target in target_order + ] + expected_targets = { + target: {"arch": arch, "multi_processor_count": multi_processor_count} + for target, (arch, multi_processor_count) in _TARGETS.items() + if target in canonical_targets + } + if ( + isinstance(catalog["schema_version"], bool) + or catalog["schema_version"] != 4 + or not canonical_targets + or target_order != canonical_targets + or not isinstance(catalog["targets"], dict) + or list(catalog["targets"]) != target_order + or catalog["targets"] != expected_targets + or catalog["domain_order"] != list(_DOMAIN_ORDER) + ): + raise RuntimeError( + "TRT-LLM MLA Blackwell generated-source catalog identity is invalid" + ) + + domains = catalog["domains"] + if not isinstance(domains, dict) or tuple(domains) != _DOMAIN_ORDER: + raise RuntimeError( + "TRT-LLM MLA Blackwell generated-source catalog domain topology is invalid" + ) + source_paths: set[str] = set() + # Every physical target carries its own generated host source and device + # identities: the schedules are traced per target, so module identities and + # host dispatch code are independent between targets. + module_idents_by_target: dict[str, set[str]] = { + target: set() for target in target_order + } + device_source_count = {target: 0 for target in target_order} + for domain, expected_device_count in _DOMAIN_DEVICE_COUNTS.items(): + profile = domains[domain] + if not isinstance(profile, dict) or set(profile) != { + "host_sources", + "device_sources", + }: + raise RuntimeError( + f"TRT-LLM MLA Blackwell catalog domain {domain!r} schema is invalid" + ) + hosts_by_target = profile["host_sources"] + devices_by_target = profile["device_sources"] + if ( + not isinstance(hosts_by_target, dict) + or list(hosts_by_target) != target_order + or not isinstance(devices_by_target, dict) + or list(devices_by_target) != target_order + ): + raise RuntimeError( + f"TRT-LLM MLA catalog domain {domain!r} target inventory is invalid" + ) + for target in target_order: + host = _source_record( + hosts_by_target[target], + domain=domain, + target=target, + kind="host_source", + ) + host_path = str(host["path"]) + if ( + host_path != f"host/{target}/cake_{domain}.cpp" + or host_path in source_paths + ): + raise RuntimeError("TRT-LLM MLA catalog host source paths are invalid") + source_paths.add(host_path) + devices = devices_by_target[target] + if not isinstance(devices, list) or len(devices) != expected_device_count: + raise RuntimeError( + f"TRT-LLM MLA catalog domain {domain!r}/{target} must contain " + f"exactly {expected_device_count} device sources" + ) + records = [ + _source_record( + device, + domain=domain, + target=target, + kind="device_source", + index=index, + ) + for index, device in enumerate(devices) + ] + idents = tuple(str(record["module_ident"]) for record in records) + module_idents = module_idents_by_target[target] + if len(set(idents)) != len(idents) or any( + ident in module_idents for ident in idents + ): + raise RuntimeError( + "TRT-LLM MLA catalog contains duplicate device module identities" + ) + module_idents.update(idents) + paths = [str(record["path"]) for record in records] + expected_paths = [f"device/{target}/cake_{ident}.cu" for ident in idents] + if paths != expected_paths or any(path in source_paths for path in paths): + raise RuntimeError( + "TRT-LLM MLA catalog device source paths are invalid" + ) + source_paths.update(paths) + device_source_count[target] += len(records) + if device_source_count != {target: 18 for target in target_order} or any( + len(module_idents_by_target[target]) != 18 for target in target_order + ): + raise RuntimeError( + "TRT-LLM MLA generated-source catalog must contain exactly 18 " + "device sources for each supported target" + ) + return catalog + + +def _domain_profile(domain: str, target: str) -> Mapping[str, object]: + if domain not in _DOMAIN_DEVICE_COUNTS: + raise ValueError( + f"unknown TRT-LLM MLA domain {domain!r}; expected one of " + f"{list(_DOMAIN_ORDER)!r}" + ) + if target not in _TARGETS: + raise ValueError( + f"unsupported TRT-LLM MLA target {target!r}; expected one of " + f"{list(_TARGET_ORDER)!r}" + ) + catalog = _source_catalog() + targets = catalog["targets"] + assert isinstance(targets, dict) + if target not in targets: + raise ValueError( + f"TRT-LLM MLA target {target!r} is absent from the generated-source " + f"catalog; available targets: {catalog['target_order']!r}" + ) + domains = catalog["domains"] + assert isinstance(domains, dict) + profile = domains[domain] + assert isinstance(profile, dict) + return profile + + +def _target_key(device: torch.device | int | str | None = None) -> str: + capability = torch.cuda.get_device_capability(device) + multi_processor_count = torch.cuda.get_device_properties( + device + ).multi_processor_count + for target, (arch, expected_multi_processor_count) in _TARGETS.items(): + if ( + capability == _ARCH_CAPABILITIES[arch] + and multi_processor_count == expected_multi_processor_count + ): + return target + raise ValueError( + "TRT-LLM MLA requires one of the exact targets " + f"{list(_TARGET_ORDER)!r}, got compute capability " + f"{capability[0]}.{capability[1]} with {multi_processor_count} SMs" + ) + + +def _sealed_source_bytes( + source_dir: Path, + record: Mapping[str, object], +) -> tuple[Path, bytes]: + relative = record["path"] + expected_sha256 = record["sha256"] + assert isinstance(relative, str) + assert isinstance(expected_sha256, str) + generated_root = (source_dir / "generated").resolve() + path = (generated_root / relative).resolve() + if generated_root not in path.parents: + raise RuntimeError( + f"TRT-LLM MLA Blackwell generated source path escapes its package: {path}" + ) + if not path.is_file(): + raise RuntimeError(f"TRT-LLM MLA Blackwell generated source is missing: {path}") + payload = path.read_bytes() + actual_sha256 = hashlib.sha256(payload).hexdigest() + if actual_sha256 != expected_sha256: + raise RuntimeError( + "TRT-LLM MLA Blackwell generated source identity drift: " + f"{path} has sha256={actual_sha256}, expected {expected_sha256}" + ) + return path, payload + + +def _nvcc() -> Path: + candidate = shutil.which("nvcc") + if candidate is None: + cuda_root = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") + if cuda_root: + path = Path(cuda_root) / "bin" / "nvcc" + if path.is_file(): + candidate = str(path) + if candidate is None: + raise RuntimeError( + "nvcc is required to build the TRT-LLM MLA Blackwell backend" + ) + return Path(candidate).resolve() + + +def _digest_field(digest, value: bytes) -> None: + digest.update(len(value).to_bytes(8, "big")) + digest.update(value) + + +@functools.cache +def _load_domain_module(domain: str, target: str): + """Compile and load one catalog-bound domain for an exact physical target.""" + + profile = _domain_profile(domain, target) + arch, multi_processor_count = _TARGETS[target] + source_dir = _source_dir() + hosts_by_target = profile["host_sources"] + devices_by_target = profile["device_sources"] + assert isinstance(hosts_by_target, dict) + assert isinstance(devices_by_target, dict) + host = hosts_by_target[target] + devices = devices_by_target[target] + assert isinstance(host, dict) + assert isinstance(devices, list) + _, host_payload = _sealed_source_bytes(source_dir, host) + nvcc = _nvcc() + + digest = hashlib.sha256() + for value in ( + domain.encode(), + target.encode(), + arch.encode(), + str(multi_processor_count).encode(), + str(nvcc).encode(), + host_payload, + ): + _digest_field(digest, value) + resolved_devices: list[tuple[str, Path, tuple[str, ...]]] = [] + for device in devices: + assert isinstance(device, dict) + module_ident = device["module_ident"] + compile_flags = device["compile_flags"] + assert isinstance(module_ident, str) + assert isinstance(compile_flags, list) + device_path, device_payload = _sealed_source_bytes(source_dir, device) + flags = tuple(compile_flags) + for value in ( + module_ident.encode(), + device_payload, + "\0".join(flags).encode(), + ): + _digest_field(digest, value) + resolved_devices.append((module_ident, device_path, flags)) + + module_name = f"trtllm_mla_{domain}_{target}_{digest.hexdigest()[:16]}" + build_dir = jit_env.FLASHINFER_JIT_DIR / module_name + build_dir.mkdir(parents=True, exist_ok=True) + lock_path = build_dir / f"{module_name}.lock" + with FileLock(lock_path, thread_local=False): + cubins: dict[str, bytes] = {} + for module_ident, device_path, compile_flags in resolved_devices: + cubin_path = build_dir / f"{module_ident}.cubin" + if not cubin_path.is_file(): + temporary = build_dir / f"{module_ident}.{os.getpid()}.tmp.cubin" + command = [ + str(nvcc), + "-cubin", + f"-arch={arch}", + "--std=c++17", + "-O3", + "-I", + str(nvcc.parent.parent / "include"), + *compile_flags, + str(device_path), + "-o", + str(temporary), + ] + process = subprocess.run(command, text=True, capture_output=True) + if process.returncode != 0: + temporary.unlink(missing_ok=True) + raise RuntimeError( + f"TRT-LLM MLA Blackwell nvcc failed for {domain}/" + f"{module_ident} ({target}, {arch}):\n{process.stderr}" + ) + os.replace(temporary, cubin_path) + cubins[module_ident] = cubin_path.read_bytes() + + module = cpp.load_inline( + module_name, + cpp_sources=host_payload.decode("utf-8"), + embed_cubin=cubins, + extra_include_paths=[str(nvcc.parent.parent / "include")], + extra_cflags=["-O3"], + extra_ldflags=["-lcuda"], + build_directory=str(build_dir), + ) + logger.info("Loaded TRT-LLM MLA domain %s for target %s (%s)", domain, target, arch) + return module + + +def get_domain_module(domain: str, device: torch.device | int | str | None = None): + """Return the cached source-built module for one exact public domain.""" + + if domain not in _DOMAIN_DEVICE_COUNTS: + raise ValueError( + f"unknown TRT-LLM MLA domain {domain!r}; expected one of " + f"{list(_DOMAIN_ORDER)!r}" + ) + return _load_domain_module(domain, _target_key(device)) + + +# --- Semantic dispatcher --- + + +_ALIGNED_HEADS = 128 +_ALIGNED_QK_DIM = 576 +_ALIGNED_VALUE_DIM = 512 +_CLUSTER_SIZE = 2 +_MAX_GENERIC_TOKENS = 4096 +_WORKSPACE_CACHE_CAPACITY = 8 +_HOST_METADATA_CACHE_CAPACITY = 16 + +_VHALF_ROUTE = "bf16_q128_runtime_vhalf_underfill_sink_one_launch_v1" +_VQUARTER_ROUTE = "bf16_q128_runtime_vquarter_underfill_sink_one_launch_v1" +_UNSPLIT_ROUTE = "bf16_b64_q16_kv1024_unsplit_v1" +_CLC_ROUTE = "bf16_clc_packed_affine_full_tile_mask_bmm2_one_v1" +_GENERIC_BF16_ROUTE = "bf16_full_abi_runtime_tail_one_launch_v1" +_GENERIC_FP8_ROUTE = "fp8_full_abi_runtime_tail_one_launch_v1" +_FP8_P32_QK_L2_ROUTE = "fp8_p32_q2_kv1024_qk_l2_resident_v1" +_FP8_PAGE64_ROUTE = "fp8_page64_native_pdl_sequence_unified_v_leader_consumer_v1" +_NATIVE_BF16_ROUTE = "v32_15stage_page_native_sink_pdl_runtime_v4_full_tmem_scrub" + +ROUTE_TO_DOMAIN = { + _VQUARTER_ROUTE: "mla_bf16_vquarter", + _VHALF_ROUTE: "mla_bf16_vhalf", + _UNSPLIT_ROUTE: "mla_bf16_unsplit", + _CLC_ROUTE: "mla_bf16_clc", + _GENERIC_BF16_ROUTE: "mla_bf16_tail", + _GENERIC_FP8_ROUTE: "mla_fp8_tail", + _FP8_P32_QK_L2_ROUTE: "mla_fp8_p32_qk_l2", + _FP8_PAGE64_ROUTE: "mla_fp8_page64_pdl", + _NATIVE_BF16_ROUTE: "mla_bf16_native_split8_pdl", +} + +_WORKSPACE_CACHE: "OrderedDict[tuple[Any, ...], dict[str, Any]]" = OrderedDict() +_WORKSPACE_CACHE_LOCK = threading.Lock() +_HOST_METADATA_CACHE: "OrderedDict[tuple[Any, ...], tuple[torch.Tensor, tuple[int, ...]]]" = OrderedDict() +_HOST_METADATA_CACHE_LOCK = threading.Lock() + + +@dataclass(frozen=True) +class _BlackwellDispatchMetadata: + """Host-visible values which determine the exported semantic domain.""" + + dtype: torch.dtype + batch_size: int + q_len: int + total_q: int + q_lens: tuple[int, ...] + kv_lens: tuple[int, ...] + num_heads: int + qk_dim: int + value_dim: int + page_size: int + max_seq_len: int + topk: int + table_ndim: int + num_sms: int + ragged_query: bool = False + uses_shared_paged_kv_idx: bool = True + enable_sink: bool = False + skip_softmax: bool = False + return_lse: bool = False + provide_lse: bool = False + device_scale: bool = False + bmm2_scale: float = 1.0 + + @property + def variant(self) -> str: + return "topk" if self.topk > 0 else "dense" + + +def _workspace_get_or_create( + key: tuple[Any, ...], builder: Callable[[], dict[str, Any]] +) -> dict[str, Any]: + with _WORKSPACE_CACHE_LOCK: + state = _WORKSPACE_CACHE.get(key) + if state is not None: + _WORKSPACE_CACHE.move_to_end(key) + return state + state = builder() + with _WORKSPACE_CACHE_LOCK: + existing = _WORKSPACE_CACHE.get(key) + if existing is not None: + _WORKSPACE_CACHE.move_to_end(key) + return existing + _WORKSPACE_CACHE[key] = state + while len(_WORKSPACE_CACHE) > _WORKSPACE_CACHE_CAPACITY: + _WORKSPACE_CACHE.popitem(last=False) + return state + + +def _tensor_ptr(value: Any) -> int: + return 0 if value is None else int(value.data_ptr()) + + +def _tensor_version(value: Any) -> int: + return -1 if value is None else int(getattr(value, "_version", 0)) + + +def _host_int_tuple(tensor: torch.Tensor) -> tuple[int, ...]: + key = ( + _tensor_ptr(tensor), + _tensor_version(tensor), + tuple(tensor.shape), + tensor.dtype, + tensor.device, + ) + with _HOST_METADATA_CACHE_LOCK: + entry = _HOST_METADATA_CACHE.get(key) + if entry is not None and entry[0] is tensor: + _HOST_METADATA_CACHE.move_to_end(key) + return entry[1] + values = tuple(int(value) for value in tensor.tolist()) + with _HOST_METADATA_CACHE_LOCK: + existing = _HOST_METADATA_CACHE.get(key) + if existing is not None and existing[0] is tensor: + _HOST_METADATA_CACHE.move_to_end(key) + return existing[1] + _HOST_METADATA_CACHE[key] = (tensor, values) + while len(_HOST_METADATA_CACHE) > _HOST_METADATA_CACHE_CAPACITY: + _HOST_METADATA_CACHE.popitem(last=False) + return values + + +def _state_key(inputs: dict[str, Any], family: str) -> tuple[Any, ...]: + tensors = ( + inputs.get("Q"), + inputs.get("KV_cache"), + inputs.get("page_table"), + inputs.get("q_indptr"), + inputs.get("seq_lens"), + inputs.get("O"), + inputs.get("LSE"), + inputs.get("sinks"), + ) + return ( + family, + *((_tensor_ptr(value), _tensor_version(value)) for value in tensors), + tuple(inputs["Q"].shape), + tuple(inputs["KV_cache"].shape), + str(inputs["dtype"]), + int(inputs["page_size"]), + int(inputs["q_len"]), + int(inputs["total_q"]), + tuple(inputs["q_lens"]), + tuple(inputs["kv_lens"]), + int(inputs["topk"]), + int(inputs["stream"]), + ) + + +def _e4m3_decode_table() -> list[float]: + values = [] + for bits in range(256): + sign = -1.0 if bits & 0x80 else 1.0 + exponent = (bits >> 3) & 0xF + mantissa = bits & 0x7 + if exponent == 0: + magnitude = mantissa * (2.0**-9) + elif exponent == 0xF and mantissa == 0x7: + magnitude = 0.0 + else: + magnitude = (1.0 + mantissa / 8.0) * (2.0 ** (exponent - 7)) + values.append(sign * magnitude) + return values + + +def _runtime_rows(inputs: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: + device = inputs["Q"].device + row_batches, causal_lens = [], [] + for batch, (q_len, kv_len) in enumerate( + zip(inputs["q_lens"], inputs["kv_lens"], strict=True) + ): + for query in range(q_len): + row_batches.append(batch) + causal_lens.append(kv_len - q_len + query + 1) + row_batches = torch.tensor(row_batches, dtype=torch.int32, device=device) + row_seq_lens = torch.tensor(causal_lens, dtype=torch.int32, device=device) + return row_batches, row_seq_lens + + +def _base_state(inputs: dict[str, Any], family: str) -> dict[str, Any]: + key = _state_key(inputs, family) + + def build() -> dict[str, Any]: + row_batches, row_seq_lens = _runtime_rows(inputs) + device = inputs["Q"].device + row_batch_long = row_batches.to(torch.int64) + use_sparse = int(inputs["topk"]) > 0 + sparse_width = int(inputs["topk"]) if use_sparse else 0 + source_table = inputs["page_table"] + if source_table.ndim == 3: + source_table = source_table[:, 0] + source_table = source_table.to(dtype=torch.int32).contiguous() + row_table = source_table.index_select(0, row_batch_long) + if use_sparse: + sparse_indices = inputs["page_table"].reshape( + int(inputs["total_q"]), sparse_width + ) + sparse_indices = sparse_indices.to(dtype=torch.int32).contiguous() + row_seq_lens = torch.full_like(row_seq_lens, sparse_width) + else: + sparse_indices = row_batches + sinks = inputs["sinks"] + if sinks is None: + sinks = torch.zeros( + int(inputs["num_heads"]), dtype=torch.float32, device=device + ) + lse = inputs["LSE"] + if lse is None: + lse = torch.empty( + (int(inputs["total_q"]), int(inputs["num_heads"])), + dtype=torch.float32, + device=device, + ) + fp8_lut = ( + torch.tensor(_e4m3_decode_table(), dtype=torch.float32, device=device) + if inputs["dtype"] == torch.float8_e4m3fn + else torch.zeros(1, dtype=torch.float32, device=device) + ) + return { + "num_rows": int(inputs["total_q"]), + "row_batches": row_batches, + "row_seq_lens": row_seq_lens, + "source_page_table": source_table, + "row_page_table": row_table, + "sparse_indices": sparse_indices, + "sparse_width": sparse_width, + "use_sparse": use_sparse, + "sinks": sinks, + "lse": lse, + "fp8_lut": fp8_lut, + } + + return _workspace_get_or_create(key, build) + + +def _longest_first_work_order(tile_counts: list[int]) -> list[int]: + """Work ids in descending KV-tile order; ties keep row order.""" + return sorted( + range(len(tile_counts)), key=lambda row: (-int(tile_counts[row]), row) + ) + + +def _append_work_order_rows( + table: torch.Tensor, inputs: dict[str, Any] +) -> torch.Tensor: + """Append the longest-first work permutation after the num_rows page-table rows. + + The persistent CLC kernels consume raw work ids in launch order and read + ``page_table[num_rows * max_pages_per_seq + raw_id]`` to find the row each id + processes, so long rows start first and the dynamic scheduler balances the + tail. Padding entries hold -1; rows below num_rows are unchanged, so every + other aligned domain sees the same page table as before. + """ + tile_counts = [ + (kv_len - q_len + query + 128) // 128 + for q_len, kv_len in zip(inputs["q_lens"], inputs["kv_lens"], strict=True) + for query in range(int(q_len)) + ] + num_rows, width = int(table.shape[0]), int(table.shape[1]) + if num_rows != len(tile_counts): + raise ValueError("aligned page table rows do not match the query rows") + order = _longest_first_work_order(tile_counts) + tail_rows = (num_rows + width - 1) // width + tail = torch.full((tail_rows * width,), -1, dtype=torch.int32, device=table.device) + tail[:num_rows] = torch.tensor(order, dtype=torch.int32, device=table.device) + return torch.cat( + (table.to(dtype=torch.int32), tail.view(tail_rows, width)), dim=0 + ).contiguous() + + +def _aligned_state(inputs: dict[str, Any]) -> dict[str, Any]: + state = _base_state(inputs, "aligned_bf16") + if "kv_half_pages" not in state: + table = state["row_page_table"] + if int(inputs["page_size"]) == 64: + table = torch.stack((table * 2, table * 2 + 1), dim=-1) + table = table.reshape(state["num_rows"], -1).contiguous() + table = _append_work_order_rows(table, inputs) + state.update( + { + "row_page_table": table, + "kv_half_pages": inputs["KV_cache"].reshape(-1, 32, _ALIGNED_QK_DIM), + "q_rows": inputs["Q"].reshape(-1, _ALIGNED_QK_DIM), + "o_rows": inputs["O"].reshape(-1, _ALIGNED_VALUE_DIM), + "num_sms": int(inputs["num_sms"]), + } + ) + return state + + +def _is_aligned(meta: _BlackwellDispatchMetadata) -> bool: + return ( + meta.num_heads == _ALIGNED_HEADS + and meta.qk_dim == _ALIGNED_QK_DIM + and meta.value_dim == _ALIGNED_VALUE_DIM + and meta.topk == 0 + and meta.uses_shared_paged_kv_idx + and not meta.return_lse + and not meta.provide_lse + ) + + +def _can_clc(meta: _BlackwellDispatchMetadata) -> bool: + q_lens = meta.q_lens + return ( + meta.dtype == torch.bfloat16 + and _is_aligned(meta) + and meta.page_size == 32 + and meta.variant == "dense" + and not meta.enable_sink + and not meta.skip_softmax + and meta.bmm2_scale == 1.0 + and bool(q_lens) + and min(q_lens) > 0 + and len(set(q_lens)) == 1 + ) + + +def _use_unsplit(meta: _BlackwellDispatchMetadata) -> bool: + return ( + _can_clc(meta) + and meta.q_lens == (16,) * 64 + and len(meta.kv_lens) == 64 + and max(meta.kv_lens) == 1024 + and meta.kv_lens[-1] == 1024 + ) + + +def _use_native_bf16(meta: _BlackwellDispatchMetadata) -> bool: + if ( + meta.dtype != torch.bfloat16 + or not _is_aligned(meta) + or meta.variant != "dense" + or meta.ragged_query + or meta.skip_softmax + ): + return False + return ( + meta.page_size == 32 + and not meta.enable_sink + and meta.q_lens == (4,) + and meta.kv_lens == (1024,) + ) or ( + meta.page_size == 64 + and meta.enable_sink + and meta.q_lens == (1,) + and meta.kv_lens == (1024,) + ) + + +def _can_fp8_p32(meta: _BlackwellDispatchMetadata) -> bool: + q_lens = meta.q_lens + if not q_lens or len(q_lens) != len(meta.kv_lens) or len(set(q_lens)) != 1: + return False + q_len = q_lens[0] + causal = [ + kv_len - q_len + query + 1 + for kv_len, query_len in zip(meta.kv_lens, q_lens, strict=True) + for query in range(query_len) + ] + return ( + meta.dtype == torch.float8_e4m3fn + and _is_aligned(meta) + and q_len == 2 + and meta.kv_lens == (1024,) + and not meta.ragged_query + and meta.page_size == 32 + and meta.variant == "dense" + and meta.table_ndim == 2 + and not meta.device_scale + and not meta.enable_sink + and not meta.skip_softmax + and all(length > 256 for length in causal) + ) + + +def _can_page64(meta: _BlackwellDispatchMetadata) -> bool: + if meta.ragged_query and len(set(meta.q_lens)) != 1: + return False + return ( + meta.dtype == torch.float8_e4m3fn + and _is_aligned(meta) + and meta.page_size == 64 + and meta.variant == "dense" + and not meta.enable_sink + and not meta.skip_softmax + ) + + +def _select_route(meta: _BlackwellDispatchMetadata) -> str: + if meta.dtype == torch.bfloat16 and _is_aligned(meta): + resident = max(1, meta.num_sms // _CLUSTER_SIZE) + if _use_native_bf16(meta): + return _NATIVE_BF16_ROUTE + if _use_unsplit(meta): + return _UNSPLIT_ROUTE + if meta.enable_sink or meta.total_q * 4 <= resident: + return _VQUARTER_ROUTE + if meta.total_q * 2 <= resident: + return _VHALF_ROUTE + if _can_clc(meta): + return _CLC_ROUTE + raise ValueError( + "TRT-LLM MLA Blackwell has no qualified aligned BF16 domain for this configuration" + ) + if _can_fp8_p32(meta): + return _FP8_P32_QK_L2_ROUTE + if _can_page64(meta): + return _FP8_PAGE64_ROUTE + if max(meta.kv_lens) > _MAX_GENERIC_TOKENS: + raise ValueError( + "TRT-LLM MLA Blackwell generic tail supports max_seq_len <= " + f"{_MAX_GENERIC_TOKENS}" + ) + return _GENERIC_BF16_ROUTE if meta.dtype == torch.bfloat16 else _GENERIC_FP8_ROUTE + + +def _clc_source_selector_eligibility( + inputs: dict[str, Any], q_lens: tuple[int, ...], kv_lens: tuple[int, ...] +) -> int: + """Resolve page-table eligibility bits 1/2 and all-KV-uniform bit 4.""" + source_page_table = inputs.get("page_table") + q4_source_eligible = ( + source_page_table is not None + and int(source_page_table.ndim) == 2 + and int(source_page_table.shape[1]) == 32 + and int(source_page_table.data_ptr()) % 16 == 0 + ) + register_profile_source_eligible = ( + source_page_table is not None + and int(source_page_table.ndim) == 2 + and int(source_page_table.shape[0]) == len(q_lens) + ) + return ( + int(q4_source_eligible) + + 2 * int(register_profile_source_eligible) + + 4 + * int( + len(kv_lens) == len(q_lens) + and all(length == kv_lens[0] for length in kv_lens) + ) + ) + + +def _aligned_launch(inputs: dict[str, Any], route: str): + state = _aligned_state(inputs) + resident = max(1, state["num_sms"] // _CLUSTER_SIZE) + if route == _VQUARTER_ROUTE: + split = 4 + work = state["num_rows"] * 4 + grid_x = min(work, resident) * 2 + grid_z = 1 + elif route == _VHALF_ROUTE: + split = 2 + work = state["num_rows"] * 2 + grid_x = min(work, resident) * 2 + grid_z = 1 + elif route == _UNSPLIT_ROUTE: + split = 1 + work = state["num_rows"] + grid_x = min(work, resident) * 2 + grid_z = 1 + else: + split = int(inputs["q_len"]) + work = state["num_rows"] + grid_x = 2 * split + grid_z = int(inputs["batch_size"]) + tensors = ( + state["q_rows"], + state["kv_half_pages"], + state["o_rows"], + state["row_seq_lens"], + state["row_page_table"], + state["sinks"], + ) + scalars: tuple[Any, ...] = ( + float(inputs["bmm1_scale"]) * log2e, + float(inputs["bmm2_scale"]), + work, + split, + int(state["row_page_table"].shape[1]), + int(inputs["sinks"] is not None), + grid_x, + grid_z, + ) + if route == _CLC_ROUTE: + kv_lens = tuple(int(value) for value in inputs["kv_lens"]) + q_lens = tuple(int(value) for value in inputs["q_lens"]) + scalars += ( + max(kv_lens), + kv_lens[-1], + _clc_source_selector_eligibility(inputs, q_lens, kv_lens), + ) + return tensors, scalars + + +def _native_launch(inputs: dict[str, Any]): + state = _aligned_state(inputs) + num_split = 8 + total_work_items = state["num_rows"] * num_split + grid_x = min(total_work_items, state["num_sms"] // _CLUSTER_SIZE) * _CLUSTER_SIZE + if "native_partial_output" not in state: + state["native_partial_output"] = torch.empty( + (state["num_rows"], _ALIGNED_HEADS, num_split, _ALIGNED_VALUE_DIM), + dtype=torch.bfloat16, + device=inputs["Q"].device, + ) + state["native_partial_lse"] = torch.empty( + (state["num_rows"], _ALIGNED_HEADS, num_split), + dtype=torch.float32, + device=inputs["Q"].device, + ) + tensors = ( + state["q_rows"], + state["kv_half_pages"], + state["o_rows"], + state["row_seq_lens"], + state["row_page_table"], + state["sinks"], + state["native_partial_output"], + state["native_partial_lse"], + ) + scalars = ( + float(inputs["bmm1_scale"]) * log2e, + 1.0, + float(inputs["bmm2_scale"]), + num_split, + total_work_items, + int(state["row_page_table"].shape[1]), + int(inputs["sinks"] is not None), + grid_x, + state["num_rows"], + ) + return tensors, scalars + + +def _generic_launch(inputs: dict[str, Any]): + state = _base_state(inputs, "generic") + num_heads = int(inputs["num_heads"]) + qk_dim = int(inputs["Q"].shape[-1]) + value_dim = int(inputs["O"].shape[-1]) + kv_stride = int(inputs["KV_cache"].shape[-1]) + query = inputs["Q"].reshape(-1, qk_dim) + kv = inputs["KV_cache"].reshape(-1, kv_stride) + if inputs["dtype"] == torch.float8_e4m3fn: + query = query.view(torch.uint8) + kv = kv.view(torch.uint8) + tensors = ( + query, + kv, + state["fp8_lut"], + state["source_page_table"], + state["sparse_indices"], + state["row_batches"], + state["row_seq_lens"], + inputs["O"].reshape(-1, value_dim), + state["lse"].reshape(-1), + state["sinks"], + ) + scalars = ( + num_heads, + qk_dim, + value_dim, + kv_stride, + int(inputs["page_size"]), + int(state["source_page_table"].shape[1]), + state["sparse_width"], + int(state["use_sparse"]), + float(inputs["bmm1_scale"]), + float(inputs["bmm2_scale"]), + int(inputs["sinks"] is not None), + int(inputs["return_lse"] or inputs["provide_lse"]), + state["num_rows"], + ) + return tensors, scalars + + +def _p32_launch(inputs: dict[str, Any]): + reduction_groups = int(inputs["batch_size"]) * 32 * int(inputs["q_len"]) + key = _state_key(inputs, "fp8_p32") + state = _workspace_get_or_create( + key, + lambda: { + "completion": torch.zeros( + reduction_groups, + dtype=torch.uint32, + device=inputs["Q"].device, + ) + }, + ) + if "partial_output" not in state: + state["partial_output"] = torch.empty( + (reduction_groups, 2, 16, 128), + dtype=torch.bfloat16, + device=inputs["Q"].device, + ) + state["partial_stats"] = torch.empty( + (reduction_groups, 2, 16, 2), + dtype=torch.float32, + device=inputs["Q"].device, + ) + tensors = ( + inputs["Q"].reshape(-1, 576).view(torch.uint8), + inputs["KV_cache"].reshape(-1, 32, 576).view(torch.uint8), + inputs["page_table"].reshape(-1), + inputs["seq_lens"], + inputs["O"].reshape(-1), + state["completion"], + state["partial_output"], + state["partial_stats"], + ) + scalars = ( + int(inputs["batch_size"]), + int(inputs["q_len"]), + int(inputs["page_table"].shape[-1]), + int(inputs["q_len"]), + 2, + float(inputs["bmm1_scale"]) * log2e, + float(inputs["bmm2_scale"]), + ) + return tensors, scalars + + +def _pick_num_split(work_items: int, tiles: list[int], num_sms: int) -> int: + tiles = [int(value) for value in tiles if int(value) > 0] + if not tiles: + return 1 + min_tiles, max_tiles = min(tiles), max(tiles) + source_cap = max(1, (max_tiles + 1) // 2) + if min_tiles <= 2: + return 1 + target = min(source_cap, max(1, num_sms // max(1, work_items * 2))) + split = max(1, min(target, min_tiles)) + blocks = (min_tiles + split - 1) // split + split = (min_tiles + blocks - 1) // blocks + while split > 1: + if all((split - 1) * ((count + split - 1) // split) < count for count in tiles): + break + split -= 1 + return split + + +def _page64_launch(inputs: dict[str, Any]): + state = _base_state(inputs, "fp8_page64") + q_lens = inputs["q_lens"] + kv_lens = inputs["kv_lens"] + work_items = sum(q_lens) + causal_lengths = [ + kv_lens[batch] - q_len + query + for batch, q_len in enumerate(q_lens) + for query in range(q_len) + ] + tiles = [(length + 127) // 128 for length in causal_lengths] + num_split = _pick_num_split(work_items, tiles, int(inputs["num_sms"])) + if num_split <= 1: + raise ValueError( + "TRT-LLM MLA Blackwell page-64 PDL domain requires split-K reduction" + ) + reduce_ctas = min( + 64, + max(1, (int(inputs["num_sms"]) * 2) // max(1, work_items * 2)), + ) + if "partial_output" not in state: + state["partial_output"] = torch.empty( + (work_items, num_split, _ALIGNED_HEADS, _ALIGNED_VALUE_DIM), + dtype=torch.bfloat16, + device=inputs["Q"].device, + ) + state["partial_stats"] = torch.empty( + (work_items, num_split, _ALIGNED_HEADS, 2), + dtype=torch.float32, + device=inputs["Q"].device, + ) + tensors = ( + inputs["Q"].reshape(-1, 576).view(torch.uint8), + inputs["KV_cache"].reshape(-1, 576).view(torch.uint8), + inputs["O"].reshape(-1, 512), + state["lse"].reshape(-1, 128), + torch.tensor(causal_lengths, dtype=torch.int32, device=inputs["Q"].device), + state["row_batches"], + inputs["page_table"].reshape(-1), + state["partial_output"], + state["partial_stats"], + ) + scalars = ( + float(inputs["bmm1_scale"]) * log2e, + float(inputs["bmm2_scale"]), + num_split, + work_items, + int(inputs["page_table"].shape[-1]), + reduce_ctas, + ) + return tensors, scalars + + +def _prepare(inputs: dict[str, Any], route: str): + if route == _NATIVE_BF16_ROUTE: + tensors, scalars = _native_launch(inputs) + elif route in {_VQUARTER_ROUTE, _VHALF_ROUTE, _UNSPLIT_ROUTE, _CLC_ROUTE}: + tensors, scalars = _aligned_launch(inputs, route) + elif route in {_GENERIC_BF16_ROUTE, _GENERIC_FP8_ROUTE}: + tensors, scalars = _generic_launch(inputs) + elif route == _FP8_P32_QK_L2_ROUTE: + tensors, scalars = _p32_launch(inputs) + elif route == _FP8_PAGE64_ROUTE: + tensors, scalars = _page64_launch(inputs) + else: + raise ValueError( + f"route is outside the TRT-LLM MLA Blackwell export inventory: {route!r}" + ) + return ROUTE_TO_DOMAIN[route], tensors, scalars + + +def _check_tensor( + tensor: torch.Tensor, + *, + name: str, + dtype: torch.dtype, + device: torch.device, +) -> None: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}") + if tensor.dtype != dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + + +def _normalize_sinks( + sinks: Optional[Union[list[torch.Tensor], tuple[torch.Tensor, ...], torch.Tensor]], + *, + num_heads: int, + device: torch.device, +) -> Optional[torch.Tensor]: + if sinks is None: + return None + if isinstance(sinks, (list, tuple)): + if len(sinks) != 1: + raise ValueError("TRT-LLM MLA Blackwell expects one sink tensor") + sink = sinks[0] + else: + sink = sinks + _check_tensor(sink, name="sinks", dtype=torch.float32, device=device) + if tuple(sink.shape) != (num_heads,): + raise ValueError(f"sinks must have shape ({num_heads},)") + return sink + + +def _normalize_scale(value: float | torch.Tensor, name: str) -> float: + if isinstance(value, torch.Tensor): + raise TypeError(f"TRT-LLM MLA Blackwell requires scalar {name}") + if not isinstance(value, (int, float)): + raise TypeError(f"{name} must be a scalar float") + result = float(value) + if not math.isfinite(result): + raise ValueError(f"{name} must be finite") + return result + + +def trtllm_mla_blackwell_decode( + query: torch.Tensor, + kv_cache: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: Optional[torch.Tensor], + max_seq_len: int, + *, + qk_nope_head_dim: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + sparse_mla_top_k: int, + out: Optional[torch.Tensor], + bmm1_scale: float | torch.Tensor, + bmm2_scale: float | torch.Tensor, + sinks: Optional[list[torch.Tensor]], + skip_softmax_threshold_scale_factor: Optional[float], + enable_pdl: Optional[bool], + uses_shared_paged_kv_idx: bool, + lse: Optional[torch.Tensor], + return_lse: bool, + cum_seq_lens_q: Optional[torch.Tensor], + max_q_len: Optional[int], + multi_ctas_kv_counter_buffer: Optional[torch.Tensor], + sparse_mla_top_k_lens: Optional[torch.Tensor], + enable_dcp: bool, + backend: Literal["cake"] = "cake", +) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + """Dispatch the qualified SM100a/SM103a MLA semantic envelope.""" + + if backend != "cake": + raise ValueError(f"backend must be 'cake', got {backend!r}") + if not isinstance(query, torch.Tensor): + raise TypeError("query must be a torch.Tensor") + if not query.is_cuda: + raise ValueError("query must be a CUDA tensor") + device = query.device + capability = get_compute_capability(device) + if capability not in {(10, 0), (10, 3)}: + major, minor = capability + raise RuntimeError( + "TRT-LLM MLA Blackwell requires compute capability 10.0 or 10.3, " + f"got {major}.{minor}" + ) + if not is_cuda_version_at_least("12.9"): + raise RuntimeError( + "TRT-LLM MLA Blackwell on SM100a/SM103a requires CUDA 12.9 or newer" + ) + if enable_dcp: + raise ValueError("TRT-LLM MLA Blackwell does not support DCP") + if multi_ctas_kv_counter_buffer is not None: + raise ValueError( + "TRT-LLM MLA Blackwell does not use a multi-CTA counter buffer" + ) + if sparse_mla_top_k_lens is not None: + raise ValueError("TRT-LLM MLA Blackwell does not accept sparse_mla_top_k_lens") + if seq_lens is None: + raise ValueError("seq_lens is required for TRT-LLM MLA Blackwell") + + if query.dtype not in {torch.bfloat16, torch.float8_e4m3fn}: + raise TypeError("TRT-LLM MLA Blackwell query must use BF16 or FP8 E4M3") + _check_tensor(query, name="query", dtype=query.dtype, device=device) + _check_tensor(kv_cache, name="kv_cache", dtype=query.dtype, device=device) + if query.ndim not in {3, 4}: + raise ValueError( + "query must be a fixed [B, Q, H, D] or compact [T, H, D] tensor" + ) + if kv_cache.ndim not in {3, 4}: + raise ValueError("kv_cache must be a 3D or 4D paged tensor") + if kv_cache.ndim == 4 and kv_cache.shape[1] != 1: + raise ValueError("4D kv_cache must have a singleton head axis") + page_size = int(kv_cache.shape[-2]) + if page_size not in {32, 64}: + raise ValueError("TRT-LLM MLA Blackwell requires page size 32 or 64") + if int(query.shape[-1]) != int(kv_cache.shape[-1]): + raise ValueError("query and kv_cache head dimensions differ") + + ragged_query = cum_seq_lens_q is not None + if ragged_query != (query.ndim == 3): + raise ValueError("compact query and cum_seq_lens_q must be provided together") + if ragged_query: + _check_tensor( + cum_seq_lens_q, + name="cum_seq_lens_q", + dtype=torch.int32, + device=device, + ) + if cum_seq_lens_q.ndim != 1 or cum_seq_lens_q.numel() < 2: + raise ValueError("cum_seq_lens_q must have shape [batch_size + 1]") + if max_q_len is None or max_q_len <= 0: + raise ValueError( + "max_q_len is required for compact TRT-LLM MLA Blackwell queries" + ) + batch_size = int(cum_seq_lens_q.numel() - 1) + q_len = int(max_q_len) + total_q = int(query.shape[0]) + else: + batch_size = int(query.shape[0]) + q_len = int(query.shape[1]) + total_q = batch_size * q_len + if batch_size <= 0 or q_len <= 0 or total_q <= 0: + raise ValueError( + "TRT-LLM MLA Blackwell requires nonempty batch and query dimensions" + ) + + _check_tensor(seq_lens, name="seq_lens", dtype=torch.int32, device=device) + if tuple(seq_lens.shape) != (batch_size,): + raise ValueError(f"seq_lens must have shape ({batch_size},)") + if ragged_query: + q_indptr_host = _host_int_tuple(cum_seq_lens_q) + if q_indptr_host[0] != 0 or q_indptr_host[-1] != total_q: + raise ValueError("cum_seq_lens_q must start at 0 and end at total_q") + q_lens = tuple( + right - left + for left, right in zip(q_indptr_host[:-1], q_indptr_host[1:], strict=True) + ) + if any(q <= 0 for q in q_lens) or max(q_lens) > q_len: + raise ValueError("cum_seq_lens_q contains an invalid query length") + else: + q_lens = (q_len,) * batch_size + kv_lens = _host_int_tuple(seq_lens) + if len(q_lens) != len(kv_lens) or any( + q <= 0 or kv <= 0 or q > kv for q, kv in zip(q_lens, kv_lens, strict=True) + ): + raise ValueError("every TRT-LLM MLA Blackwell row requires 0 < q_len <= kv_len") + if max_seq_len <= 0 or max(kv_lens) > max_seq_len: + raise ValueError("max_seq_len must cover every runtime KV length") + _check_tensor(block_tables, name="block_tables", dtype=torch.int32, device=device) + if sparse_mla_top_k > 0: + expected_table_shape = ( + (total_q, sparse_mla_top_k) + if ragged_query + else (batch_size, q_len, sparse_mla_top_k) + ) + if tuple(block_tables.shape) != expected_table_shape: + raise ValueError( + f"sparse block_tables must have shape {expected_table_shape}" + ) + else: + expected_ndim = 2 if uses_shared_paged_kv_idx else 3 + if block_tables.ndim != expected_ndim or block_tables.shape[0] != batch_size: + raise ValueError( + "dense block_tables must match the batch and shared-index layout" + ) + + num_heads = int(query.shape[-2]) + qk_dim = int(query.shape[-1]) + value_dim = int(kv_lora_rank) + valid_dense = sparse_mla_top_k == 0 and ( + qk_nope_head_dim, + kv_lora_rank, + qk_rope_head_dim, + num_heads, + ) in { + (128, 512, 64, 128), + (128, 512, 64, 64), + (64, 256, 64, 32), + (512, 512, 64, 128), + } + valid_topk = sparse_mla_top_k > 0 and ( + qk_nope_head_dim, + kv_lora_rank, + qk_rope_head_dim, + num_heads, + ) in { + (128, 512, 64, 128), + (128, 512, 64, 64), + (192, 512, 64, 128), + (192, 512, 64, 64), + } + if not valid_dense and not valid_topk: + raise ValueError("unsupported TRT-LLM MLA Blackwell dimension tuple") + if qk_dim != kv_lora_rank + qk_rope_head_dim: + raise ValueError("query width must equal kv_lora_rank + qk_rope_head_dim") + + sink = _normalize_sinks(sinks, num_heads=num_heads, device=device) + bmm1_scale_value = _normalize_scale(bmm1_scale, "bmm1_scale") + bmm2_scale_value = _normalize_scale(bmm2_scale, "bmm2_scale") + if skip_softmax_threshold_scale_factor is not None: + threshold = float(skip_softmax_threshold_scale_factor) + if not math.isfinite(threshold) or threshold <= 0: + raise ValueError("skip_softmax_threshold_scale_factor must be positive") + expected_out_shape = (*query.shape[:-1], value_dim) + if out is None: + out = torch.empty(expected_out_shape, dtype=torch.bfloat16, device=device) + else: + _check_tensor(out, name="out", dtype=torch.bfloat16, device=device) + if tuple(out.shape) != expected_out_shape: + raise ValueError(f"out must have shape {expected_out_shape}") + if lse is not None: + _check_tensor(lse, name="lse", dtype=torch.float32, device=device) + valid_lse_shapes = {(total_q, num_heads), (*query.shape[:-1],)} + if tuple(lse.shape) not in valid_lse_shapes: + raise ValueError("lse shape must match flattened or physical query rows") + if return_lse and lse is None: + lse = torch.empty((total_q, num_heads), dtype=torch.float32, device=device) + + num_sms = get_device_sm_count(device) + metadata = _BlackwellDispatchMetadata( + dtype=query.dtype, + batch_size=batch_size, + q_len=q_len, + total_q=total_q, + q_lens=q_lens, + kv_lens=kv_lens, + num_heads=num_heads, + qk_dim=qk_dim, + value_dim=value_dim, + page_size=page_size, + max_seq_len=int(max_seq_len), + topk=int(sparse_mla_top_k), + table_ndim=int(block_tables.ndim), + num_sms=num_sms, + ragged_query=ragged_query, + uses_shared_paged_kv_idx=uses_shared_paged_kv_idx, + enable_sink=sink is not None, + skip_softmax=skip_softmax_threshold_scale_factor is not None, + return_lse=return_lse, + provide_lse=lse is not None and not return_lse, + device_scale=False, + bmm2_scale=bmm2_scale_value, + ) + route = _select_route(metadata) + stream = int(torch.cuda.current_stream(device).cuda_stream) + inputs = { + "Q": query, + "KV_cache": kv_cache, + "page_table": block_tables, + "q_indptr": cum_seq_lens_q, + "seq_lens": seq_lens, + "O": out, + "LSE": lse, + "sinks": sink, + "dtype": query.dtype, + "batch_size": batch_size, + "q_len": q_len, + "total_q": total_q, + "q_lens": q_lens, + "kv_lens": kv_lens, + "num_heads": num_heads, + "page_size": page_size, + "max_seq_len": int(max_seq_len), + "topk": int(sparse_mla_top_k), + "bmm1_scale": bmm1_scale_value, + "bmm2_scale": bmm2_scale_value, + "return_lse": return_lse, + "provide_lse": lse is not None and not return_lse, + "num_sms": num_sms, + "stream": stream, + } + domain, tensors, scalars = _prepare(inputs, route) + with torch.cuda.device(device): + get_domain_module(domain, device).run(*tensors, *scalars, stream) + if return_lse: + assert lse is not None + return out, lse + return out + + +__all__ = ["ROUTE_TO_DOMAIN", "trtllm_mla_blackwell_decode"] diff --git a/tests/attention/test_trtllm_gen_mla.py b/tests/attention/test_trtllm_gen_mla.py index f66e3b8a226..0498bf1c0ad 100755 --- a/tests/attention/test_trtllm_gen_mla.py +++ b/tests/attention/test_trtllm_gen_mla.py @@ -1017,6 +1017,44 @@ def test_trtllm_batch_decode_mla_non_power_of_two_heads( ) +def trtllm_mla_blackwell_reference( + query: torch.Tensor, + kv_cache: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + kv_flat = kv_cache.reshape(-1, 576) + outputs = [] + for batch_index in range(query.shape[0]): + seq_len = int(seq_lens[batch_index].item()) + pages = block_tables[batch_index, : (seq_len + 31) // 32] + indices = ( + pages[:, None] * 32 + + torch.arange(32, device=query.device, dtype=torch.int32)[None, :] + ).reshape(-1)[:seq_len] + batch_outputs = [] + for query_index in range(query.shape[1]): + right = seq_len - query.shape[1] + query_index + kv = kv_flat[indices.long()][: right + 1] + logits = torch.einsum( + "hd,kd->hk", + query[batch_index, query_index, :, :512].float(), + kv[:, :512].float(), + ) + logits += torch.einsum( + "hd,kd->hk", + query[batch_index, query_index, :, 512:].float(), + kv[:, 512:].float(), + ) + probabilities = F.softmax(logits * softmax_scale, dim=-1) + batch_outputs.append( + torch.einsum("hk,kd->hd", probabilities, kv[:, :512].float()) + ) + outputs.append(torch.stack(batch_outputs)) + return torch.stack(outputs).to(torch.bfloat16) + + @pytest.mark.parametrize( "num_heads,dtype,max_seq_len", [ @@ -1815,6 +1853,179 @@ def test_trtllm_mla_prefill_matches_decode_multi_token_bf16(): torch.testing.assert_close(tensor_only_result, decode_out, rtol=1e-2, atol=1e-2) +@pytest.mark.parametrize( + "overrides,expected_domain", + [ + ( + {"batch_size": 1, "q_len": 4, "total_q": 4}, + "mla_bf16_native_split8_pdl", + ), + ( + { + "batch_size": 1, + "q_len": 1, + "total_q": 1, + "page_size": 64, + "enable_sink": True, + }, + "mla_bf16_native_split8_pdl", + ), + ( + {"batch_size": 1, "q_len": 8, "total_q": 8}, + "mla_bf16_vquarter", + ), + ( + {"batch_size": 4, "q_len": 8, "total_q": 32}, + "mla_bf16_vhalf", + ), + ( + {"batch_size": 4, "q_len": 16, "total_q": 64}, + "mla_bf16_clc", + ), + ( + {"batch_size": 64, "q_len": 16, "total_q": 1024}, + "mla_bf16_unsplit", + ), + ( + {"num_heads": 32, "qk_dim": 320, "value_dim": 256}, + "mla_bf16_tail", + ), + ( + { + "dtype": torch.float8_e4m3fn, + "num_heads": 32, + "qk_dim": 320, + "value_dim": 256, + }, + "mla_fp8_tail", + ), + ( + {"dtype": torch.float8_e4m3fn, "q_len": 2, "total_q": 2}, + "mla_fp8_p32_qk_l2", + ), + ( + { + "dtype": torch.float8_e4m3fn, + "q_len": 16, + "total_q": 16, + "page_size": 64, + "max_seq_len": 4096, + }, + "mla_fp8_page64_pdl", + ), + ], +) +def test_trtllm_mla_blackwell_semantic_domain_selection( + overrides: dict, expected_domain: str +) -> None: + from flashinfer.mla.cake_trtllm_mla_blackwell import ( + ROUTE_TO_DOMAIN, + _BlackwellDispatchMetadata, + _select_route, + ) + + values = { + "dtype": torch.bfloat16, + "batch_size": 1, + "q_len": 8, + "total_q": 8, + "num_heads": 128, + "qk_dim": 576, + "value_dim": 512, + "page_size": 32, + "max_seq_len": 1024, + "topk": 0, + "table_ndim": 2, + "num_sms": 148, + } + values.update(overrides) + values["q_lens"] = (values["q_len"],) * values["batch_size"] + values["kv_lens"] = (values["max_seq_len"],) * values["batch_size"] + route = _select_route(_BlackwellDispatchMetadata(**values)) + assert ROUTE_TO_DOMAIN[route] == expected_domain + + +def test_trtllm_mla_blackwell_rejects_non_scalar_or_nonfinite_scales() -> None: + from flashinfer.mla.cake_trtllm_mla_blackwell import _normalize_scale + + with pytest.raises(TypeError, match="requires scalar bmm1_scale"): + _normalize_scale(torch.ones(1), "bmm1_scale") + with pytest.raises(ValueError, match="bmm2_scale must be finite"): + _normalize_scale(float("inf"), "bmm2_scale") + + +def test_trtllm_mla_blackwell_caches_host_lengths_until_tensor_version_changes() -> ( + None +): + from flashinfer.mla.cake_trtllm_mla_blackwell import ( + _HOST_METADATA_CACHE, + _host_int_tuple, + ) + + _HOST_METADATA_CACHE.clear() + lengths = torch.tensor([1024, 768], dtype=torch.int32) + first = _host_int_tuple(lengths) + assert _host_int_tuple(lengths) is first + lengths.add_(1) + assert _host_int_tuple(lengths) == (1025, 769) + + +@pytest.mark.parametrize( + "batch_size,q_len_per_request", + [(1, 4), (1, 8), (4, 2), (4, 4), (4, 8), (4, 16)], +) +def test_trtllm_mla_blackwell_bf16_dispatch( + batch_size: int, q_len_per_request: int +) -> None: + if not torch.cuda.is_available(): + pytest.skip("TRT-LLM MLA Blackwell requires CUDA") + if get_compute_capability(torch.device("cuda")) not in {(10, 0), (10, 3)}: + pytest.skip("TRT-LLM MLA Blackwell requires SM100a or SM103a") + + torch.manual_seed(42) + device = torch.device("cuda") + page_size = 32 + max_seq_len = 1024 + pages_per_sequence = max_seq_len // page_size + query = ( + torch.randn(batch_size, q_len_per_request, 128, 576, device=device) * 0.05 + ).to(torch.bfloat16) + kv_cache = ( + torch.randn(batch_size * pages_per_sequence, page_size, 576, device=device) + * 0.05 + ).to(torch.bfloat16) + block_tables = torch.arange( + batch_size * pages_per_sequence, dtype=torch.int32, device=device + ).reshape(batch_size, pages_per_sequence) + seq_lens = torch.full((batch_size,), max_seq_len, dtype=torch.int32, device=device) + workspace = torch.empty(1, dtype=torch.uint8, device=device) + bmm1_scale = 1.0 / (192**0.5) + + output = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( + query=query, + kv_cache=kv_cache, + workspace_buffer=workspace, + qk_nope_head_dim=128, + kv_lora_rank=512, + qk_rope_head_dim=64, + block_tables=block_tables, + seq_lens=seq_lens, + max_seq_len=max_seq_len, + bmm1_scale=bmm1_scale, + bmm2_scale=1.0, + backend="cake", + ) + reference = trtllm_mla_blackwell_reference( + query, + kv_cache, + block_tables, + seq_lens, + bmm1_scale, + ) + assert output.shape == (batch_size, q_len_per_request, 128, 512) + torch.testing.assert_close(output, reference, rtol=1e-2, atol=1e-2) + + @pytest.mark.parametrize( "layer_dimensions", supported_mla_layer_dimensions,