From 05ce5994aef5803b1556cba805511bea9a0d2c22 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Mon, 2 Nov 2026 02:50:20 -0800 Subject: [PATCH 1/2] perf(cake_rmsnorm_det): faster H 256 dw reduce and dc main, register-compensated H 1536 backward Three backward changes, each bitwise equal (dx and dw) to the program it replaces: - H = 256 dw reduce: a column-split form of the slab-16 chain (one column of a four-column slab per lane, the same chains in the same order) runs over 64 CTAs instead of 16. Its TwoSum form reduces the compensated-dot windows on SM100 (above 80896 rows), SM103 (above 99840) and SM107 (above 70656); its Kahan form the SM103 ring window between 40960 and 99840 rows. - H = 256 compensated-dot ring on SM100 and SM107: a lean form (two rows per iteration, the same arithmetic in the same row order) serves its windows. - H = 1536 above the accuracy edge on SM103 (40448 rows) and SM107 (45056 rows): the compensated block sums are kept in registers instead of the shared-memory fold buffer. The programs they replace stay in the package. cake_windows.py records the measured class of the backward edges measured since, and the pairs test follows the lean ring. --- .../cake_rmsnorm_det_train_backend/build.json | 276 ++++++ .../cake_jit.py | 361 +++++++- .../cake_windows.py | 79 +- ..._det_train_083ea823c691e0631ad3_binding.cu | 137 +++ ...m_det_train_083ea823c691e0631ad3_kernel.cu | 768 ++++++++++++++++ ..._det_train_0f0b782f1e538caad3ce_binding.cu | 146 +++ ...m_det_train_0f0b782f1e538caad3ce_kernel.cu | 832 ++++++++++++++++++ ..._det_train_16a5e0b1afd90f38ed44_binding.cu | 146 +++ ...m_det_train_16a5e0b1afd90f38ed44_kernel.cu | 813 +++++++++++++++++ ..._det_train_affa1b95f68bab8c411b_binding.cu | 80 ++ ...m_det_train_affa1b95f68bab8c411b_kernel.cu | 178 ++++ ..._det_train_b4b63d5b5028c7a60333_binding.cu | 80 ++ ...m_det_train_b4b63d5b5028c7a60333_kernel.cu | 199 +++++ ..._det_train_c4c89f0d04b7b3a38d0b_binding.cu | 137 +++ ...m_det_train_c4c89f0d04b7b3a38d0b_kernel.cu | 780 ++++++++++++++++ ..._det_train_ca6211f3ea2ce8c86a05_binding.cu | 146 +++ ...m_det_train_ca6211f3ea2ce8c86a05_kernel.cu | 819 +++++++++++++++++ ..._det_train_e694b0a25ea01d1f6cdc_binding.cu | 146 +++ ...m_det_train_e694b0a25ea01d1f6cdc_kernel.cu | 826 +++++++++++++++++ .../norm/test_cake_rmsnorm_det_train_pairs.py | 18 +- 20 files changed, 6898 insertions(+), 69 deletions(-) create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_binding.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_kernel.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_binding.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_kernel.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_binding.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_kernel.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_affa1b95f68bab8c411b_binding.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_affa1b95f68bab8c411b_kernel.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_binding.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_kernel.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_binding.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_kernel.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_binding.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_kernel.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_binding.cu create mode 100644 flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_kernel.cu diff --git a/flashinfer/cake_rmsnorm_det_train_backend/build.json b/flashinfer/cake_rmsnorm_det_train_backend/build.json index e81a30ce6f6..6309afdffb6 100644 --- a/flashinfer/cake_rmsnorm_det_train_backend/build.json +++ b/flashinfer/cake_rmsnorm_det_train_backend/build.json @@ -382,6 +382,52 @@ "csrc/cake_rmsnorm_det_train_07f9f2fbea7163639c27_binding.cu" ] }, + { + "arch": "sm_100a", + "closure_sha256": "191b54c66f29b2825918b0b46d33295fe2e0f6e7a9fb7e523f0fb8ee0a1f1338", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_rmsnorm_det_train_083ea823c691e0631ad3_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_rmsnorm_det_train_083ea823c691e0631ad3_sm_100a.so", + "sha256": { + "csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_binding.cu": "ca99d17ed79fc6da7a79159f63a5e24a94f582ed25eeea94419785188fd38014", + "csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_kernel.cu": "8429d64278ee1addb327ee2794980268b98573ef2207d642f6113afe7d9aec63" + }, + "sources": [ + "csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_kernel.cu", + "csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "471c6984cdea8fcb42a480a868d07c91981a93541cd209b83cf33238ed5de8d4", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_rmsnorm_det_train_083ea823c691e0631ad3_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_rmsnorm_det_train_083ea823c691e0631ad3_sm_107a.so", + "sha256": { + "csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_binding.cu": "ca99d17ed79fc6da7a79159f63a5e24a94f582ed25eeea94419785188fd38014", + "csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_kernel.cu": "8429d64278ee1addb327ee2794980268b98573ef2207d642f6113afe7d9aec63" + }, + "sources": [ + "csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_kernel.cu", + "csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_binding.cu" + ] + }, { "arch": "sm_100a", "closure_sha256": "4621ecb13b0842962c0793627548e5fb59782efaea9779f2914c0a1038164a3e", @@ -750,6 +796,29 @@ "csrc/cake_rmsnorm_det_train_0dcdc563be377247044e_binding.cu" ] }, + { + "arch": "sm_107a", + "closure_sha256": "de4e9f3f086070d51369c2aa344fd56cfb20364a75638d223f6ef6f3bb78cb12", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_rmsnorm_det_train_0f0b782f1e538caad3ce_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_sm_107a.so", + "sha256": { + "csrc/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_binding.cu": "1bd8486a494d7036afbe16eb11c21863c594a5af563ae822e922183a87b1d8b8", + "csrc/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_kernel.cu": "37fcce0d271c8a639012b95ee426968f59b134bd10cd1b94394a56de9f74c72a" + }, + "sources": [ + "csrc/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_kernel.cu", + "csrc/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_binding.cu" + ] + }, { "arch": "sm_107a", "closure_sha256": "c4e55c69b08bbe4a88c166aba2e20876347b413469a967cb6fb71db182ffe155", @@ -957,6 +1026,29 @@ "csrc/cake_rmsnorm_det_train_16545056022785f2a7d8_binding.cu" ] }, + { + "arch": "sm_103a", + "closure_sha256": "c572f35f026dd73157f07b75a1886721de4c07c1e676853e49bece5a6e53aa68", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_sm_103a.so", + "sha256": { + "csrc/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_binding.cu": "d166d71faebff5c2af3a2125cdfdccebe5c514f23c0433628ce04ddb33a86a15", + "csrc/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_kernel.cu": "081416008770c6de4469667a9b71938158dd64c419e115fb07a85bbfe015b062" + }, + "sources": [ + "csrc/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_kernel.cu", + "csrc/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_binding.cu" + ] + }, { "arch": "sm_100a", "closure_sha256": "ae148571acfeb8c4e4c9396e74a0b9a48c9022b2270e2578b6327b7ea3240901", @@ -6086,6 +6178,29 @@ "csrc/cake_rmsnorm_det_train_af005462cd6a301f31c4_binding.cu" ] }, + { + "arch": "sm_103a", + "closure_sha256": "6e18bed98fa4c97b7ad782c690c10617336d370563636ef21e9724d11447e35a", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_rmsnorm_det_train_affa1b95f68bab8c411b_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_rmsnorm_det_train_affa1b95f68bab8c411b_sm_103a.so", + "sha256": { + "csrc/cake_rmsnorm_det_train_affa1b95f68bab8c411b_binding.cu": "435b176f050fec35c007aadb649d390930a9e746d8a9b95e8d9778824075679f", + "csrc/cake_rmsnorm_det_train_affa1b95f68bab8c411b_kernel.cu": "f1b15fd14ae040ed00e55ef1275315a55c39669b8aeb608490415844e9533f3a" + }, + "sources": [ + "csrc/cake_rmsnorm_det_train_affa1b95f68bab8c411b_kernel.cu", + "csrc/cake_rmsnorm_det_train_affa1b95f68bab8c411b_binding.cu" + ] + }, { "arch": "sm_100a", "closure_sha256": "f5c0e7f4638b8c73fe59b3e6cdec18c0c0296d2b41105e720a2b38b2c60d4556", @@ -6339,6 +6454,75 @@ "csrc/cake_rmsnorm_det_train_b49de9b47372fa4e1051_binding.cu" ] }, + { + "arch": "sm_100a", + "closure_sha256": "edd139c506006ad1eceacfcfb07d3b32fbba45be261914ec41080c7e9a37e18f", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_rmsnorm_det_train_b4b63d5b5028c7a60333_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_sm_100a.so", + "sha256": { + "csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_binding.cu": "6e914437ef4c04143041f5c27eee5f4da274049c114ce49717d1463bcacebfbc", + "csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_kernel.cu": "e2b25a16b0ad0739265a8a5558aadcb38be9c954a0a328c8560b03966a62e67f" + }, + "sources": [ + "csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_kernel.cu", + "csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "b8cae5a135c191f93c62ed0f146dfadae043e7f77834c4ce9b4afc603c0e6215", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_rmsnorm_det_train_b4b63d5b5028c7a60333_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_sm_103a.so", + "sha256": { + "csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_binding.cu": "6e914437ef4c04143041f5c27eee5f4da274049c114ce49717d1463bcacebfbc", + "csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_kernel.cu": "e2b25a16b0ad0739265a8a5558aadcb38be9c954a0a328c8560b03966a62e67f" + }, + "sources": [ + "csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_kernel.cu", + "csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "39bfa69faeca68ecfd8974216f19c87fa6b848460b7572fec3890442ec505295", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_rmsnorm_det_train_b4b63d5b5028c7a60333_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_sm_107a.so", + "sha256": { + "csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_binding.cu": "6e914437ef4c04143041f5c27eee5f4da274049c114ce49717d1463bcacebfbc", + "csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_kernel.cu": "e2b25a16b0ad0739265a8a5558aadcb38be9c954a0a328c8560b03966a62e67f" + }, + "sources": [ + "csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_kernel.cu", + "csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_binding.cu" + ] + }, { "arch": "sm_100a", "closure_sha256": "45ef2bb49d2e17a3289fc86a35ee5a52bccaece9090eb453e4a5947e0d9392be", @@ -7167,6 +7351,52 @@ "csrc/cake_rmsnorm_det_train_c440e57b9cd55c3e22e9_binding.cu" ] }, + { + "arch": "sm_100a", + "closure_sha256": "39fe82893f54d7ddd3cb156e6d4f56cab6b3feb6c0f579814116085b701cef30", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_sm_100a.so", + "sha256": { + "csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_binding.cu": "0810b730497661b41c0bd8f73f5efd29c21461dbdf19b81cd5006a4e608076c0", + "csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_kernel.cu": "f5fc4e0824e2ba07bc1abbabc4ef2f5b5dfdbcad214504445ff8847f2e9b9bcb" + }, + "sources": [ + "csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_kernel.cu", + "csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "c6ef952c1327816191c5b93d0e93ce6b81514ee7bf927e339470e8aeb763aaf1", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_sm_107a.so", + "sha256": { + "csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_binding.cu": "0810b730497661b41c0bd8f73f5efd29c21461dbdf19b81cd5006a4e608076c0", + "csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_kernel.cu": "f5fc4e0824e2ba07bc1abbabc4ef2f5b5dfdbcad214504445ff8847f2e9b9bcb" + }, + "sources": [ + "csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_kernel.cu", + "csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_binding.cu" + ] + }, { "arch": "sm_107a", "closure_sha256": "5cd4970bdb0febd0b1fadecde22127b4a7ebef81fda48f6f4602964e4e05c50c", @@ -7351,6 +7581,29 @@ "csrc/cake_rmsnorm_det_train_c9f55c0ae4b198c859d1_binding.cu" ] }, + { + "arch": "sm_103a", + "closure_sha256": "308e2acf18f9132fd2bbb9e3f32cc3106a6a5fe30d4da132ed5b4763db4c435c", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_sm_103a.so", + "sha256": { + "csrc/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_binding.cu": "c4860ef5f927b3eb5549c653b6f535871eae4534b791531f56532eb67099d7d3", + "csrc/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_kernel.cu": "df1b2c013e8e26d33b567c2f9582fc81e76ea9b2c0d02df1e7b991a0be5ddb96" + }, + "sources": [ + "csrc/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_kernel.cu", + "csrc/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_binding.cu" + ] + }, { "arch": "sm_103a", "closure_sha256": "2f5349ea663e8d0d3811409220f56c342152f1e119a5f36af72f940c40449abc", @@ -8179,6 +8432,29 @@ "csrc/cake_rmsnorm_det_train_e646bbdbd852884953da_binding.cu" ] }, + { + "arch": "sm_107a", + "closure_sha256": "1b6b0467ecfd4a83fe8d8447247c01d4d8eedaadc37fe6ea74c4237926f04718", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_sm_107a.so", + "sha256": { + "csrc/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_binding.cu": "ef315f6c62a34f2c13bb9f905f198db6ae8f6e69221c9e5f6f9cfac9b68b1fa7", + "csrc/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_kernel.cu": "26727e817d14493d2d028ad780fbfc3ab50d064ce0c14bb0aa3dc393b18f1dbc" + }, + "sources": [ + "csrc/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_kernel.cu", + "csrc/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_binding.cu" + ] + }, { "arch": "sm_100a", "closure_sha256": "cb0997480f794dce6744a4a96f73ea6f64a081dcdbe1a9ac5d11a36caec235d1", diff --git a/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py b/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py index 1eb6e689d01..c06fea652d5 100644 --- a/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py +++ b/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py @@ -5012,6 +5012,69 @@ "kind": "chunk_partials", }, }, + "cake_rmsnorm_det_train_16a5e0b1afd90f38ed44": { + "arches": [ + "sm_103a", + ], + "grid_rule": { + "ctas_per_sm": 1, + "grid_factor": 2, + "kind": "claimed_chunks", + "launch_bound_ctas_per_sm": 1, + "max_halving": 5, + "rows_per_issue": 4, + "tail_div": 12, + }, + "hidden": 1536, + "kernel": "cake_rmsnorm_det_train_16a5e0b1afd90f38ed44", + "launch": { + "cooperative": False, + "use_pdl": False, + }, + "library": "cake_rmsnorm_det_train_16a5e0b1afd90f38ed44", + "role": "bwd_large", + "select_rule": { + "above": 40448, + "at_most": None, + "kind": "rows_between", + }, + "workspace_rule": { + "counters": 48, + "kind": "chunk_partials", + }, + }, + # The FP32-weight twin of cake_rmsnorm_det_train_16a5e0b1afd90f38ed44 (the effective weight read at FP32; every other operand and the host plan the owner's). + "cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05": { + "arches": [ + "sm_103a", + ], + "grid_rule": { + "ctas_per_sm": 1, + "grid_factor": 2, + "kind": "claimed_chunks", + "launch_bound_ctas_per_sm": 1, + "max_halving": 5, + "rows_per_issue": 4, + "tail_div": 12, + }, + "hidden": 1536, + "kernel": "cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05", + "launch": { + "cooperative": False, + "use_pdl": False, + }, + "library": "cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05", + "role": "bwd_large_w32", + "select_rule": { + "above": 40448, + "at_most": None, + "kind": "rows_between", + }, + "workspace_rule": { + "counters": 48, + "kind": "chunk_partials", + }, + }, "cake_rmsnorm_det_train_1863320f3158d35eead6": { "arches": [ "sm_100a", @@ -6244,6 +6307,69 @@ "kind": "chunk_partials", }, }, + "cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc": { + "arches": [ + "sm_107a", + ], + "grid_rule": { + "ctas_per_sm": 1, + "grid_factor": 2, + "kind": "claimed_chunks", + "launch_bound_ctas_per_sm": 1, + "max_halving": 5, + "rows_per_issue": 4, + "tail_div": 12, + }, + "hidden": 1536, + "kernel": "cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc", + "launch": { + "cooperative": False, + "use_pdl": False, + }, + "library": "cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc", + "role": "bwd_main", + "select_rule": { + "above": 45056, + "at_most": None, + "kind": "rows_between", + }, + "workspace_rule": { + "counters": 48, + "kind": "chunk_partials", + }, + }, + # The FP32-weight twin of cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc (the effective weight read at FP32; every other operand and the host plan the owner's). + "cake_rmsnorm_det_train_0f0b782f1e538caad3ce": { + "arches": [ + "sm_107a", + ], + "grid_rule": { + "ctas_per_sm": 1, + "grid_factor": 2, + "kind": "claimed_chunks", + "launch_bound_ctas_per_sm": 1, + "max_halving": 5, + "rows_per_issue": 4, + "tail_div": 12, + }, + "hidden": 1536, + "kernel": "cake_rmsnorm_det_train_0f0b782f1e538caad3ce", + "launch": { + "cooperative": False, + "use_pdl": False, + }, + "library": "cake_rmsnorm_det_train_0f0b782f1e538caad3ce", + "role": "bwd_main_w32", + "select_rule": { + "above": 45056, + "at_most": None, + "kind": "rows_between", + }, + "workspace_rule": { + "counters": 48, + "kind": "chunk_partials", + }, + }, "cake_rmsnorm_det_train_6a428da2227972ae18ae": { "arches": [ "sm_107a", @@ -9679,6 +9805,69 @@ "kind": "chunk_partials", }, }, + "cake_rmsnorm_det_train_083ea823c691e0631ad3": { + "arches": [ + "sm_100a", + ], + "grid_rule": { + "ctas_per_sm": 5, + "grid_factor": 5, + "kind": "claimed_chunks", + "launch_bound_ctas_per_sm": 5, + "max_halving": 5, + "rows_per_issue": 4, + "tail_div": 4, + }, + "hidden": 256, + "kernel": "cake_rmsnorm_det_train_083ea823c691e0631ad3", + "launch": { + "cooperative": False, + "use_pdl": False, + }, + "library": "cake_rmsnorm_det_train_083ea823c691e0631ad3", + "role": "bwd_main", + "select_rule": { + "above": 80896, + "at_most": 1572864, + "kind": "rows_between", + }, + "workspace_rule": { + "counters": 128, + "kind": "chunk_partials", + }, + }, + # The FP32-weight twin of cake_rmsnorm_det_train_083ea823c691e0631ad3 (the effective weight read at FP32; every other operand and the host plan the owner's). + "cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b": { + "arches": [ + "sm_100a", + ], + "grid_rule": { + "ctas_per_sm": 5, + "grid_factor": 5, + "kind": "claimed_chunks", + "launch_bound_ctas_per_sm": 5, + "max_halving": 5, + "rows_per_issue": 4, + "tail_div": 4, + }, + "hidden": 256, + "kernel": "cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b", + "launch": { + "cooperative": False, + "use_pdl": False, + }, + "library": "cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b", + "role": "bwd_main_w32", + "select_rule": { + "above": 80896, + "at_most": 1572864, + "kind": "rows_between", + }, + "workspace_rule": { + "counters": 128, + "kind": "chunk_partials", + }, + }, "cake_rmsnorm_det_train_0291949ba535c646c3a8__99840_1572864": { "arches": [ "sm_103a", @@ -9805,6 +9994,69 @@ "kind": "chunk_partials", }, }, + "cake_rmsnorm_det_train_083ea823c691e0631ad3__70656_1572864": { + "arches": [ + "sm_107a", + ], + "grid_rule": { + "ctas_per_sm": 5, + "grid_factor": 5, + "kind": "claimed_chunks", + "launch_bound_ctas_per_sm": 5, + "max_halving": 5, + "rows_per_issue": 4, + "tail_div": 4, + }, + "hidden": 256, + "kernel": "cake_rmsnorm_det_train_083ea823c691e0631ad3", + "launch": { + "cooperative": False, + "use_pdl": False, + }, + "library": "cake_rmsnorm_det_train_083ea823c691e0631ad3", + "role": "bwd_main", + "select_rule": { + "above": 70656, + "at_most": 1572864, + "kind": "rows_between", + }, + "workspace_rule": { + "counters": 128, + "kind": "chunk_partials", + }, + }, + # The FP32-weight twin of cake_rmsnorm_det_train_083ea823c691e0631ad3__70656_1572864 (the effective weight read at FP32; every other operand and the host plan the owner's). + "cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b__70656_1572864": { + "arches": [ + "sm_107a", + ], + "grid_rule": { + "ctas_per_sm": 5, + "grid_factor": 5, + "kind": "claimed_chunks", + "launch_bound_ctas_per_sm": 5, + "max_halving": 5, + "rows_per_issue": 4, + "tail_div": 4, + }, + "hidden": 256, + "kernel": "cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b", + "launch": { + "cooperative": False, + "use_pdl": False, + }, + "library": "cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b", + "role": "bwd_main_w32", + "select_rule": { + "above": 70656, + "at_most": 1572864, + "kind": "rows_between", + }, + "workspace_rule": { + "counters": 128, + "kind": "chunk_partials", + }, + }, "cake_rmsnorm_det_train_c18297d75c0d424978fd": { "arches": [ "sm_100a", @@ -9993,21 +10245,21 @@ }, "workspace_rule": None, }, - "cake_rmsnorm_det_train_72138dce40d43da44aaf__40960_99840": { + "cake_rmsnorm_det_train_affa1b95f68bab8c411b": { "arches": [ "sm_103a", ], "grid_rule": { - "ctas": 16, + "ctas": 64, "kind": "fixed", }, "hidden": 256, - "kernel": "cake_rmsnorm_det_train_72138dce40d43da44aaf", + "kernel": "cake_rmsnorm_det_train_affa1b95f68bab8c411b", "launch": { "cooperative": False, "use_pdl": True, }, - "library": "cake_rmsnorm_det_train_72138dce40d43da44aaf", + "library": "cake_rmsnorm_det_train_affa1b95f68bab8c411b", "role": "bwd_reduce", "select_rule": { "above": 40960, @@ -10085,6 +10337,29 @@ }, "workspace_rule": None, }, + "cake_rmsnorm_det_train_b4b63d5b5028c7a60333": { + "arches": [ + "sm_100a", + ], + "grid_rule": { + "ctas": 64, + "kind": "fixed", + }, + "hidden": 256, + "kernel": "cake_rmsnorm_det_train_b4b63d5b5028c7a60333", + "launch": { + "cooperative": False, + "use_pdl": True, + }, + "library": "cake_rmsnorm_det_train_b4b63d5b5028c7a60333", + "role": "bwd_reduce", + "select_rule": { + "above": 80896, + "at_most": 1572864, + "kind": "rows_between", + }, + "workspace_rule": None, + }, "cake_rmsnorm_det_train_0ac243f8385c32960e14__99840_1572864": { "arches": [ "sm_103a", @@ -10108,6 +10383,29 @@ }, "workspace_rule": None, }, + "cake_rmsnorm_det_train_b4b63d5b5028c7a60333__99840_1572864": { + "arches": [ + "sm_103a", + ], + "grid_rule": { + "ctas": 64, + "kind": "fixed", + }, + "hidden": 256, + "kernel": "cake_rmsnorm_det_train_b4b63d5b5028c7a60333", + "launch": { + "cooperative": False, + "use_pdl": True, + }, + "library": "cake_rmsnorm_det_train_b4b63d5b5028c7a60333", + "role": "bwd_reduce", + "select_rule": { + "above": 99840, + "at_most": 1572864, + "kind": "rows_between", + }, + "workspace_rule": None, + }, "cake_rmsnorm_det_train_b30c75390d923b0168ca": { "arches": [ "sm_107a", @@ -10131,6 +10429,29 @@ }, "workspace_rule": None, }, + "cake_rmsnorm_det_train_b4b63d5b5028c7a60333__70656_1572864": { + "arches": [ + "sm_107a", + ], + "grid_rule": { + "ctas": 64, + "kind": "fixed", + }, + "hidden": 256, + "kernel": "cake_rmsnorm_det_train_b4b63d5b5028c7a60333", + "launch": { + "cooperative": False, + "use_pdl": True, + }, + "library": "cake_rmsnorm_det_train_b4b63d5b5028c7a60333", + "role": "bwd_reduce", + "select_rule": { + "above": 70656, + "at_most": 1572864, + "kind": "rows_between", + }, + "workspace_rule": None, + }, "cake_rmsnorm_det_train_e73321cec89d33336ea1": { "arches": [ "sm_100a", @@ -11575,14 +11896,14 @@ "cake_rmsnorm_det_train_a385ae1681a2c078344e__5120_12288", "cake_rmsnorm_det_train_af005462cd6a301f31c4__12288_40960", "cake_rmsnorm_det_train_a385ae1681a2c078344e__40960_80896", - "cake_rmsnorm_det_train_0291949ba535c646c3a8", + "cake_rmsnorm_det_train_083ea823c691e0631ad3", "cake_rmsnorm_det_train_c18297d75c0d424978fd", "cake_rmsnorm_det_train_72138dce40d43da44aaf", "cake_rmsnorm_det_train_72138dce40d43da44aaf__2048_5120", "cake_rmsnorm_det_train_72138dce40d43da44aaf__5120_12288", "cake_rmsnorm_det_train_72138dce40d43da44aaf__12288_40960", "cake_rmsnorm_det_train_72138dce40d43da44aaf__40960_80896", - "cake_rmsnorm_det_train_0ac243f8385c32960e14", + "cake_rmsnorm_det_train_b4b63d5b5028c7a60333", "cake_rmsnorm_det_train_e73321cec89d33336ea1", ], "bwd_w32_h256": [ @@ -11591,14 +11912,14 @@ "cake_rmsnorm_det_train_420788e663a2c3a810f6__5120_12288", "cake_rmsnorm_det_train_b7d9577ff884de98ea06__12288_40960", "cake_rmsnorm_det_train_420788e663a2c3a810f6__40960_80896", - "cake_rmsnorm_det_train_89505762fa1cbf33cb8f", + "cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b", "cake_rmsnorm_det_train_f788ef57163552825d27", "cake_rmsnorm_det_train_72138dce40d43da44aaf", "cake_rmsnorm_det_train_72138dce40d43da44aaf__2048_5120", "cake_rmsnorm_det_train_72138dce40d43da44aaf__5120_12288", "cake_rmsnorm_det_train_72138dce40d43da44aaf__12288_40960", "cake_rmsnorm_det_train_72138dce40d43da44aaf__40960_80896", - "cake_rmsnorm_det_train_0ac243f8385c32960e14", + "cake_rmsnorm_det_train_b4b63d5b5028c7a60333", "cake_rmsnorm_det_train_e73321cec89d33336ea1", ], "bwd_h4096": [ @@ -11795,7 +12116,7 @@ "cake_rmsnorm_det_train_1df7db17277d363b5168", "cake_rmsnorm_det_train_4df08ce51fb01ba1193f", "cake_rmsnorm_det_train_f522f5b8d4343d96ca04", - "cake_rmsnorm_det_train_cbeae168d70edd196417", + "cake_rmsnorm_det_train_16a5e0b1afd90f38ed44", "cake_rmsnorm_det_train_9b727c2c9933b2ca6724", "cake_rmsnorm_det_train_1863320f3158d35eead6", "cake_rmsnorm_det_train_9b727c2c9933b2ca6724__4096_14928", @@ -11809,7 +12130,7 @@ "cake_rmsnorm_det_train_a99e2b380514c4d86d60", "cake_rmsnorm_det_train_7a72e97bd6cb4dc0efc2", "cake_rmsnorm_det_train_109f19bb177c8f20cadc", - "cake_rmsnorm_det_train_6601baab5c1c5f1847e7", + "cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05", "cake_rmsnorm_det_train_9b727c2c9933b2ca6724", "cake_rmsnorm_det_train_1863320f3158d35eead6", "cake_rmsnorm_det_train_9b727c2c9933b2ca6724__4096_14928", @@ -11857,8 +12178,8 @@ "cake_rmsnorm_det_train_72138dce40d43da44aaf__2048_5120", "cake_rmsnorm_det_train_72138dce40d43da44aaf__5120_12288", "cake_rmsnorm_det_train_72138dce40d43da44aaf__12288_40960", - "cake_rmsnorm_det_train_72138dce40d43da44aaf__40960_99840", - "cake_rmsnorm_det_train_0ac243f8385c32960e14__99840_1572864", + "cake_rmsnorm_det_train_affa1b95f68bab8c411b", + "cake_rmsnorm_det_train_b4b63d5b5028c7a60333__99840_1572864", "cake_rmsnorm_det_train_e73321cec89d33336ea1", ], "bwd_w32_h256": [ @@ -11873,8 +12194,8 @@ "cake_rmsnorm_det_train_72138dce40d43da44aaf__2048_5120", "cake_rmsnorm_det_train_72138dce40d43da44aaf__5120_12288", "cake_rmsnorm_det_train_72138dce40d43da44aaf__12288_40960", - "cake_rmsnorm_det_train_72138dce40d43da44aaf__40960_99840", - "cake_rmsnorm_det_train_0ac243f8385c32960e14__99840_1572864", + "cake_rmsnorm_det_train_affa1b95f68bab8c411b", + "cake_rmsnorm_det_train_b4b63d5b5028c7a60333__99840_1572864", "cake_rmsnorm_det_train_e73321cec89d33336ea1", ], "bwd_h4096": [ @@ -12062,7 +12383,7 @@ "cake_rmsnorm_det_train_7dc29629ff21625607f8__107", "cake_rmsnorm_det_train_6a428da2227972ae18ae", "cake_rmsnorm_det_train_56d26996239835327501", - "cake_rmsnorm_det_train_bb711cc7dfaa3b1fb20c", + "cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc", "cake_rmsnorm_det_train_628248febbdfe17c3325", "cake_rmsnorm_det_train_db5c0b673c516bd0295d", "cake_rmsnorm_det_train_db5c0b673c516bd0295d__16384_45056", @@ -12072,7 +12393,7 @@ "cake_rmsnorm_det_train_baff9ecbafe7de1e6030__107", "cake_rmsnorm_det_train_e63beae98c236ce26335", "cake_rmsnorm_det_train_de8e07d4cb312b06dbc4", - "cake_rmsnorm_det_train_ec11f5e3e3fca820cb45", + "cake_rmsnorm_det_train_0f0b782f1e538caad3ce", "cake_rmsnorm_det_train_628248febbdfe17c3325", "cake_rmsnorm_det_train_db5c0b673c516bd0295d", "cake_rmsnorm_det_train_db5c0b673c516bd0295d__16384_45056", @@ -12105,21 +12426,21 @@ "bwd_h256": [ "cake_rmsnorm_det_train_af005462cd6a301f31c4__0_65536", "cake_rmsnorm_det_train_a385ae1681a2c078344e__65536_70656", - "cake_rmsnorm_det_train_0291949ba535c646c3a8__70656_1572864", + "cake_rmsnorm_det_train_083ea823c691e0631ad3__70656_1572864", "cake_rmsnorm_det_train_c18297d75c0d424978fd", "cake_rmsnorm_det_train_312bbe084eac69d858c4", "cake_rmsnorm_det_train_312bbe084eac69d858c4__65536_70656", - "cake_rmsnorm_det_train_b30c75390d923b0168ca", + "cake_rmsnorm_det_train_b4b63d5b5028c7a60333__70656_1572864", "cake_rmsnorm_det_train_158c887660637d5533d0", ], "bwd_w32_h256": [ "cake_rmsnorm_det_train_b7d9577ff884de98ea06__0_65536", "cake_rmsnorm_det_train_420788e663a2c3a810f6__65536_70656", - "cake_rmsnorm_det_train_89505762fa1cbf33cb8f__70656_1572864", + "cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b__70656_1572864", "cake_rmsnorm_det_train_f788ef57163552825d27", "cake_rmsnorm_det_train_312bbe084eac69d858c4", "cake_rmsnorm_det_train_312bbe084eac69d858c4__65536_70656", - "cake_rmsnorm_det_train_b30c75390d923b0168ca", + "cake_rmsnorm_det_train_b4b63d5b5028c7a60333__70656_1572864", "cake_rmsnorm_det_train_158c887660637d5533d0", ], "bwd_h4096": [ diff --git a/flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py b/flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py index baf4af89d0a..d5b9dca63f3 100644 --- a/flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py +++ b/flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py @@ -65,7 +65,7 @@ (5120, 12288, "cake_rmsnorm_det_train_a385ae1681a2c078344e"), (12288, 40960, "cake_rmsnorm_det_train_af005462cd6a301f31c4"), (40960, 80896, "cake_rmsnorm_det_train_a385ae1681a2c078344e"), - (80896, 1572864, "cake_rmsnorm_det_train_0291949ba535c646c3a8"), + (80896, 1572864, "cake_rmsnorm_det_train_083ea823c691e0631ad3"), (1572864, None, "cake_rmsnorm_det_train_c18297d75c0d424978fd"), ), "bwd_w32_h256": ( @@ -74,7 +74,7 @@ (5120, 12288, "cake_rmsnorm_det_train_420788e663a2c3a810f6"), (12288, 40960, "cake_rmsnorm_det_train_b7d9577ff884de98ea06"), (40960, 80896, "cake_rmsnorm_det_train_420788e663a2c3a810f6"), - (80896, 1572864, "cake_rmsnorm_det_train_89505762fa1cbf33cb8f"), + (80896, 1572864, "cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b"), (1572864, None, "cake_rmsnorm_det_train_f788ef57163552825d27"), ), "bwd_h4096": ( @@ -237,7 +237,7 @@ (4096, 14928, "cake_rmsnorm_det_train_1df7db17277d363b5168"), (14928, 15566, "cake_rmsnorm_det_train_4df08ce51fb01ba1193f"), (15566, 40448, "cake_rmsnorm_det_train_f522f5b8d4343d96ca04"), - (40448, None, "cake_rmsnorm_det_train_cbeae168d70edd196417"), + (40448, None, "cake_rmsnorm_det_train_16a5e0b1afd90f38ed44"), ), "bwd_w32_h1536": ( (0, 1024, "cake_rmsnorm_det_train_baff9ecbafe7de1e6030"), @@ -245,7 +245,7 @@ (4096, 14928, "cake_rmsnorm_det_train_a99e2b380514c4d86d60"), (14928, 15566, "cake_rmsnorm_det_train_7a72e97bd6cb4dc0efc2"), (15566, 40448, "cake_rmsnorm_det_train_109f19bb177c8f20cadc"), - (40448, None, "cake_rmsnorm_det_train_6601baab5c1c5f1847e7"), + (40448, None, "cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05"), ), "bwd_h2048": ( (0, 3072, "cake_rmsnorm_det_train_caf6978ed0a17fbf65b9"), @@ -436,13 +436,13 @@ (0, 1024, "cake_rmsnorm_det_train_7dc29629ff21625607f8"), (1024, 16384, "cake_rmsnorm_det_train_6a428da2227972ae18ae"), (16384, 45056, "cake_rmsnorm_det_train_56d26996239835327501"), - (45056, None, "cake_rmsnorm_det_train_bb711cc7dfaa3b1fb20c"), + (45056, None, "cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc"), ), "bwd_w32_h1536": ( (0, 1024, "cake_rmsnorm_det_train_baff9ecbafe7de1e6030"), (1024, 16384, "cake_rmsnorm_det_train_e63beae98c236ce26335"), (16384, 45056, "cake_rmsnorm_det_train_de8e07d4cb312b06dbc4"), - (45056, None, "cake_rmsnorm_det_train_ec11f5e3e3fca820cb45"), + (45056, None, "cake_rmsnorm_det_train_0f0b782f1e538caad3ce"), ), "bwd_h2048": ( (0, 256, "cake_rmsnorm_det_train_84efa2af7907877dfb56"), @@ -461,13 +461,13 @@ "bwd_h256": ( (0, 65536, "cake_rmsnorm_det_train_af005462cd6a301f31c4"), (65536, 70656, "cake_rmsnorm_det_train_a385ae1681a2c078344e"), - (70656, 1572864, "cake_rmsnorm_det_train_0291949ba535c646c3a8"), + (70656, 1572864, "cake_rmsnorm_det_train_083ea823c691e0631ad3"), (1572864, None, "cake_rmsnorm_det_train_c18297d75c0d424978fd"), ), "bwd_w32_h256": ( (0, 65536, "cake_rmsnorm_det_train_b7d9577ff884de98ea06"), (65536, 70656, "cake_rmsnorm_det_train_420788e663a2c3a810f6"), - (70656, 1572864, "cake_rmsnorm_det_train_89505762fa1cbf33cb8f"), + (70656, 1572864, "cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b"), (1572864, None, "cake_rmsnorm_det_train_f788ef57163552825d27"), ), "bwd_h4096": ( @@ -646,33 +646,28 @@ (5120, "same"), (12288, "same"), (40960, "same"), - (80896, "unmeasured"), - (1572864, "unmeasured"), + (80896, "same"), + (1572864, "same"), ), "bwd_w32_h256": ( (2048, "same"), (5120, "same"), (12288, "same"), (40960, "same"), - (80896, "unmeasured"), - (1572864, "unmeasured"), - ), - "bwd_h4096": ( - (256, "same"), - (1024, "same"), - (4096, "same"), - (44032, "unmeasured"), + (80896, "same"), + (1572864, "same"), ), + "bwd_h4096": ((256, "same"), (1024, "same"), (4096, "same"), (44032, "same")), "bwd_w32_h4096": ( (256, "same"), (1024, "same"), (4096, "same"), - (44032, "unmeasured"), + (44032, "same"), ), "bwd_h512": ((1024, "ulp"), (4096, "same"), (32768, "same")), "bwd_w32_h512": ((1024, "ulp"), (4096, "same"), (32768, "same")), - "bwd_h5120": ((256, "unmeasured"), (6144, "unmeasured")), - "bwd_w32_h5120": ((256, "unmeasured"), (6144, "unmeasured")), + "bwd_h5120": ((256, "same"), (6144, "same")), + "bwd_w32_h5120": ((256, "same"), (6144, "same")), "bwd_h6144": ((256, "same"), (4096, "same"), (42496, "same")), "bwd_w32_h6144": ((256, "same"), (4096, "same"), (42496, "same")), "bwd_residual_h6144": (), @@ -712,15 +707,15 @@ "bwd_h1536": ( (1024, "same"), (4096, "same"), - (14928, "unmeasured"), - (15566, "unmeasured"), + (14928, "same"), + (15566, "same"), (40448, "same"), ), "bwd_w32_h1536": ( (1024, "same"), (4096, "same"), - (14928, "unmeasured"), - (15566, "unmeasured"), + (14928, "same"), + (15566, "same"), (40448, "same"), ), "bwd_h2048": ( @@ -742,16 +737,16 @@ (5120, "same"), (12288, "same"), (40960, "same"), - (99840, "unmeasured"), - (1572864, "unmeasured"), + (99840, "same"), + (1572864, "same"), ), "bwd_w32_h256": ( (2048, "same"), (5120, "same"), (12288, "same"), (40960, "same"), - (99840, "unmeasured"), - (1572864, "unmeasured"), + (99840, "same"), + (1572864, "same"), ), "bwd_h4096": ((256, "same"), (1024, "same"), (4096, "same"), (48640, "same")), "bwd_w32_h4096": ( @@ -762,8 +757,8 @@ ), "bwd_h512": ((1024, "ulp"), (4096, "same"), (32768, "same")), "bwd_w32_h512": ((1024, "ulp"), (4096, "same"), (32768, "same")), - "bwd_h5120": ((5120, "unmeasured"),), - "bwd_w32_h5120": ((5120, "unmeasured"),), + "bwd_h5120": ((5120, "same"),), + "bwd_w32_h5120": ((5120, "same"),), "bwd_h6144": ((4096, "same"), (48640, "same")), "bwd_w32_h6144": ((4096, "same"), (48640, "same")), "bwd_residual_h6144": (), @@ -1149,40 +1144,40 @@ * ``sm_100a`` ``bwd_h2048``: (0, 3072] caf6978e = (3072, 4096] 9ac8d21e ~ (4096, 16384] c73074cf = (16384, 58880] ceeef043 = (58880, inf) 0cb817d1 * ``sm_100a`` ``bwd_h256``: (0, 2048] a385ae16 = (2048, 5120] af005462 = - (5120, 12288] a385ae16 = (12288, 40960] af005462 = (40960, 80896] a385ae16 ? - (80896, 1572864] 0291949b ? (1572864, inf) c18297d7 + (5120, 12288] a385ae16 = (12288, 40960] af005462 = (40960, 80896] a385ae16 = + (80896, 1572864] 083ea823 = (1572864, inf) c18297d7 * ``sm_100a`` ``bwd_h4096``: (0, 256] 3e644ea1 = (256, 1024] 59270fb2 = - (1024, 4096] 9a830956 = (4096, 44032] 35a0b0ca ? (44032, inf) 7122dbcd + (1024, 4096] 9a830956 = (4096, 44032] 35a0b0ca = (44032, inf) 7122dbcd * ``sm_100a`` ``bwd_h512``: (0, 1024] 09601d42 ~ (1024, 4096] 123db0b2 = (4096, 32768] ad1a5a3c = (32768, inf) 8877322a -* ``sm_100a`` ``bwd_h5120``: (0, 256] db01cec1 ? (256, 6144] b482c84d ? +* ``sm_100a`` ``bwd_h5120``: (0, 256] db01cec1 = (256, 6144] b482c84d = (6144, inf) c145ef19 * ``sm_100a`` ``bwd_h6144``: (0, 256] b93bf060 = (256, 4096] 5c804aed = (4096, 42496] b93bf060 = (42496, inf) c440e57b * ``sm_100a`` ``bwd_residual_h6144``: (0, inf) 0dcdc563 * ``sm_103a`` ``bwd_h1536``: (0, 1024] 7dc29629 = (1024, 4096] c1a9efdf = - (4096, 14928] 1df7db17 ? (14928, 15566] 4df08ce5 ? (15566, 40448] f522f5b8 = - (40448, inf) cbeae168 + (4096, 14928] 1df7db17 = (14928, 15566] 4df08ce5 = (15566, 40448] f522f5b8 = + (40448, inf) 16a5e0b1 * ``sm_103a`` ``bwd_h2048``: (0, 3072] caf6978e = (3072, 4096] 9ac8d21e ~ (4096, 16384] caf6978e = (16384, 32768] caf6978e = (32768, 55296] 302f64db = (55296, inf) 018f0cd7 * ``sm_103a`` ``bwd_h256``: (0, 2048] a385ae16 = (2048, 5120] af005462 = - (5120, 12288] a385ae16 = (12288, 40960] af005462 = (40960, 99840] a385ae16 ? - (99840, 1572864] 0291949b ? (1572864, inf) c18297d7 + (5120, 12288] a385ae16 = (12288, 40960] af005462 = (40960, 99840] a385ae16 = + (99840, 1572864] 0291949b = (1572864, inf) c18297d7 * ``sm_103a`` ``bwd_h4096``: (0, 256] 50044ee3 = (256, 1024] ad2bb26e = (1024, 4096] 24643edc = (4096, 48640] 35a0b0ca = (48640, inf) 850eec3a * ``sm_103a`` ``bwd_h512``: (0, 1024] 09601d42 ~ (1024, 4096] 8e7f4d81 = (4096, 32768] 274b0055 = (32768, inf) cfaae3ce -* ``sm_103a`` ``bwd_h5120``: (0, 5120] 2fc4932f ? (5120, inf) 7ffb443a +* ``sm_103a`` ``bwd_h5120``: (0, 5120] 2fc4932f = (5120, inf) 7ffb443a * ``sm_103a`` ``bwd_h6144``: (0, 4096] 5c804aed = (4096, 48640] 322574af = (48640, inf) 289dc9ab * ``sm_103a`` ``bwd_residual_h6144``: (0, inf) 0dcdc563 * ``sm_107a`` ``bwd_h1536``: (0, 1024] 7dc29629 = (1024, 16384] 6a428da2 = - (16384, 45056] 56d26996 = (45056, inf) bb711cc7 + (16384, 45056] 56d26996 = (45056, inf) e694b0a2 * ``sm_107a`` ``bwd_h2048``: (0, 256] 84efa2af = (256, 1024] 845cdcd8 ~ (1024, 16384] 81c3d8ab = (16384, 50688] 845cdcd8 = (50688, inf) 694decfb * ``sm_107a`` ``bwd_h256``: (0, 65536] af005462 ? (65536, 70656] a385ae16 ? - (70656, 1572864] 0291949b ? (1572864, inf) c18297d7 + (70656, 1572864] 083ea823 ? (1572864, inf) c18297d7 * ``sm_107a`` ``bwd_h4096``: (0, 256] 0da5dd4b = (256, 1024] a647fccc = (1024, 4096] d41a514f = (4096, 24064] ee4a9053 = (24064, inf) 99a01f2a * ``sm_107a`` ``bwd_h512``: (0, 1024] 3feed502 = (1024, 4096] 3feed502 = diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_binding.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_binding.cu new file mode 100644 index 00000000000..480c0a62ccd --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_binding.cu @@ -0,0 +1,137 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_det_train_083ea823c691e0631ad3, for the device source that defines it. Do not edit. +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_rmsnorm_det_train_083ea823c691e0631ad3(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0); + +namespace cake_binding_kernel_cake_rmsnorm_det_train_083ea823c691e0631ad3 { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_det_train_083ea823c691e0631ad3"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckRank(const TensorView& t, const char* name, int rank) { + TVM_FFI_CHECK(t.ndim() >= rank, ValueError) + << KERNEL_NAME << ": '" << name << "' needs at least " << rank << " dims, got " << t.ndim(); +} + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +void Run(TensorView a_g, TensorView a_x, TensorView a_w, TensorView a_r, TensorView a_dx, TensorView a_dw, TensorView a_partial, TensorView a_counters, TensorView a_g_h, int64_t a_T, int64_t a_n_chunks, int64_t a_rows_per_chunk, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_g.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_g, "g"); + check_dtype(a_g, DLDataType{kDLBfloat, 16, 1}, "g"); + check_cuda_tensor(a_x, "x"); + check_dtype(a_x, DLDataType{kDLBfloat, 16, 1}, "x"); + check_cuda_tensor(a_w, "w"); + check_dtype(a_w, DLDataType{kDLBfloat, 16, 1}, "w"); + check_cuda_tensor(a_r, "r"); + check_dtype(a_r, DLDataType{kDLFloat, 32, 1}, "r"); + check_cuda_tensor(a_dx, "dx"); + check_dtype(a_dx, DLDataType{kDLBfloat, 16, 1}, "dx"); + check_cuda_tensor(a_dw, "dw"); + check_dtype(a_dw, DLDataType{kDLFloat, 32, 1}, "dw"); + check_cuda_tensor(a_partial, "partial"); + check_dtype(a_partial, DLDataType{kDLFloat, 32, 1}, "partial"); + check_cuda_tensor(a_counters, "counters"); + check_dtype(a_counters, DLDataType{kDLUInt, 32, 1}, "counters"); + check_cuda_tensor(a_g_h, "g_h"); + check_dtype(a_g_h, DLDataType{kDLBfloat, 16, 1}, "g_h"); + TVM_FFI_CHECK(a_T >= -2147483648LL && a_T <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'T' = " << a_T + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_n_chunks >= -2147483648LL && a_n_chunks <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'n_chunks' = " << a_n_chunks + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_rows_per_chunk >= -2147483648LL && a_rows_per_chunk <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'rows_per_chunk' = " << a_rows_per_chunk + << " is outside the i32 range [-2147483648, 2147483647]"; + CheckRank(a_g, "g", 1); + CheckUnitInnerStride(a_g, "g"); + CheckRank(a_x, "x", 1); + CheckUnitInnerStride(a_x, "x"); + check_same_device(a_x, a_g, "x", "g"); + check_contiguous(a_w, "w"); + check_same_device(a_w, a_g, "w", "g"); + check_contiguous(a_r, "r"); + check_same_device(a_r, a_g, "r", "g"); + check_contiguous(a_dx, "dx"); + check_same_device(a_dx, a_g, "dx", "g"); + check_contiguous(a_dw, "dw"); + check_same_device(a_dw, a_g, "dw", "g"); + check_contiguous(a_partial, "partial"); + check_same_device(a_partial, a_g, "partial", "g"); + check_contiguous(a_counters, "counters"); + check_same_device(a_counters, a_g, "counters", "g"); + CheckRank(a_g_h, "g_h", 1); + CheckUnitInnerStride(a_g_h, "g_h"); + check_same_device(a_g_h, a_g, "g_h", "g"); + TVM_FFI_CHECK((a_n_chunks >= 1LL), ValueError) + << KERNEL_NAME << ": requires (n_chunks >= 1)"; + TVM_FFI_CHECK((a_rows_per_chunk >= 1LL), ValueError) + << KERNEL_NAME << ": requires (rows_per_chunk >= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + void* v_g = DataPtr(a_g); + void* v_x = DataPtr(a_x); + void* v_w = DataPtr(a_w); + void* v_r = DataPtr(a_r); + void* v_dx = DataPtr(a_dx); + void* v_dw = DataPtr(a_dw); + void* v_partial = DataPtr(a_partial); + void* v_counters = DataPtr(a_counters); + void* v_g_h = DataPtr(a_g_h); + int32_t v_T = static_cast(a_T); + int32_t v_n_chunks = static_cast(a_n_chunks); + int32_t v_rows_per_chunk = static_cast(a_rows_per_chunk); + int64_t v_g_stride_p0 = a_g.stride(0); + int64_t v_x_stride_p0 = a_x.stride(0); + int64_t v_g_h_stride_p0 = a_g_h.stride(0); + void* kernel_args[] = {&v_g, &v_x, &v_w, &v_r, &v_dx, &v_dw, &v_partial, &v_counters, &v_g_h, &v_T, &v_n_chunks, &v_rows_per_chunk, &v_g_stride_p0, &v_x_stride_p0, &v_g_h_stride_p0}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(128, 1, 1); + config.dynamicSmemBytes = 43648; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_rmsnorm_det_train_083ea823c691e0631ad3), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_rmsnorm_det_train_083ea823c691e0631ad3 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_det_train_083ea823c691e0631ad3::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_kernel.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_kernel.cu new file mode 100644 index 00000000000..39bd49caa70 --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_083ea823c691e0631ad3_kernel.cu @@ -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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define SMEM_DOT_SLOTS_OFF 0 +#define SMEM_DOT_SLOTS_STAGE_BYTES 32 +#define SMEM_DOT_SLOTS_STRIDE 32 +#define SMEM_FLAGS_OFF 128 +#define SMEM_FLAGS_STAGE_BYTES 32 +#define SMEM_FLAGS_STRIDE 32 +#define SMEM_RED_FOLD_OFF 256 +#define SMEM_RED_FOLD_STAGE_BYTES 2048 +#define SMEM_RED_FOLD_STRIDE 2048 +#define SMEM_ROW_FOLD_OFF 2304 +#define SMEM_ROW_FOLD_STAGE_BYTES 4096 +#define SMEM_ROW_FOLD_STRIDE 4096 +#define SMEM_COMP_SMEM_OFF 6400 +#define SMEM_COMP_SMEM_STAGE_BYTES 128 +#define SMEM_COMP_SMEM_STRIDE 128 +#define SMEM_W_SM_OFF 6528 +#define SMEM_W_SM_STAGE_BYTES 16 +#define SMEM_W_SM_STRIDE 16 +#define SMEM_PART_SM_OFF 6656 +#define SMEM_PART_SM_STAGE_BYTES 16 +#define SMEM_PART_SM_STRIDE 16 +#define SMEM_RING_OFF 6784 +#define SMEM_RING_STAGE_BYTES 36864 +#define SMEM_RING_STRIDE 36864 +#define SMEM_TOTAL 43648 +#define THREADS 128 +#define LAUNCH_MIN_BLOCKS 5 + +#include + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.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 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; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__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; +} + +extern "C" { + +__global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void +kernel_cake_rmsnorm_det_train_083ea823c691e0631ad3(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + float* dot_slots = reinterpret_cast(smem_raw + 0); + const int dot_slots_addr = smem + 0; + int* flags = reinterpret_cast(smem_raw + 128); + const int flags_addr = smem + 128; + float* red_fold = reinterpret_cast(smem_raw + 256); + const int red_fold_addr = smem + 256; + float* row_fold = reinterpret_cast(smem_raw + 2304); + const int row_fold_addr = smem + 2304; + float* comp_smem = reinterpret_cast(smem_raw + 6400); + const int comp_smem_addr = smem + 6400; + unsigned int* w_sm = reinterpret_cast(smem_raw + 6528); + const int w_sm_addr = smem + 6528; + float* part_sm = reinterpret_cast(smem_raw + 6656); + const int part_sm_addr = smem + 6656; + uint8_t* ring = reinterpret_cast(smem_raw + 6784); + const int ring_addr = smem + 6784; + + // === Task calls (declaration order) === + asm volatile("griddepcontrol.wait;" ::: "memory"); + unsigned int nxt = 0; + if (tid == 0) { + unsigned int _atomic_fetch_add_0 = atomicAdd(&counters[0], 1); + nxt = _atomic_fetch_add_0; + } + long long g_stride = g_stride_p0; + long long x_stride = x_stride_p0; + long long gh_stride = g_h_stride_p0; + int c = 0; + int rg = warp; + int tir = tid - rg * 32; + int row_begin = 0; + int row_end = 0; + int n_rows = 0; + int n_iters = 0; + int last_row = 0; + float wc[8]; + unsigned int ww[4]; + #pragma unroll + for (int v = 0; v < 1; v++) { + int wcolp = (v * 32 + tir) * 8; + { + const uint4* _vptr_1 = reinterpret_cast(w + (unsigned long long)wcolp); + uint4* _vdst_1 = reinterpret_cast(&ww[v * 4]); + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vdst_1[_blk] = _vptr_1[_blk]; + } + } + } + float acc[8]; + float comp[8]; + unsigned int compw[4]; + if (tid == 0) { + flags[1] = (int)nxt; + } + __syncthreads(); + c = flags[1]; + #pragma unroll 1 + for (int claim_it = 0; claim_it < n_chunks; claim_it++) { + if (c >= n_chunks) { + break; + } + if (tid == 0) { + unsigned int _atomic_fetch_add_2 = atomicAdd(&counters[0], 1); + nxt = _atomic_fetch_add_2; + } + int stages = (T + 3) / 4; + int sl = rows_per_chunk; + if (sl > stages) { + sl = stages; + } + if (sl < 1) { + sl = 1; + } + int share = (stages + sl - 1) / sl; + int tail = (share + 3) / 4; + if (tail < 1) { + tail = 1; + } + int rem = stages; + int size = (share + 1) / 2; + int first = 0; + int base = 0; + int rows = 0; + int halving = 1; + int begin_it = 0; + #pragma unroll 1 + for (int k = 0; k < 5; k++) { + if (halving != 0) { + if (size > tail) { + if (rem >= sl * size) { + if (c < first + sl) { + rows = size; + begin_it = base + (c - first); + halving = 0; + } else { + first = first + sl; + base = base + sl * size; + rem = rem - sl * size; + size = (size + 1) / 2; + } + } else { + halving = 0; + } + } else { + halving = 0; + } + } + } + if (rows == 0) { + int quot = rem / sl; + int extra = rem - quot * sl; + if (c < first + extra) { + rows = quot + 1; + begin_it = base + (c - first); + } else if (c < first + sl) { + rows = quot; + begin_it = base + (c - first); + } else { + begin_it = stages; + } + } + row_begin = begin_it * 4; + row_end = row_begin + rows * 4; + if (row_begin > T) { + row_begin = T; + } + if (row_end > T) { + row_end = T; + } + n_rows = row_end - row_begin; + n_iters = (n_rows + 3) / 4; + last_row = row_end - 1; + int rstep = sl * 4; + row_end = T; + n_iters = rows; + last_row = T - 1; + #pragma unroll + for (int j = 0; j < 8; j++) { + acc[j] = 0.0f; + comp[j] = 0.0f; + } + unsigned int gwA[4]; + unsigned int xwA[4]; + unsigned int ghwA[4]; + float rrA[1]; + unsigned int gwB[4]; + unsigned int xwB[4]; + unsigned int ghwB[4]; + float rrB[1]; + if (n_iters > 0) { + int row0 = ((row_begin + rg) < (last_row) ? (row_begin + rg) : (last_row)); + unsigned int sg = (unsigned int)ring_addr + (unsigned int)(rg * 9216) + (unsigned int)(tir * 16); + unsigned int sr = (unsigned int)ring_addr + (unsigned int)(rg * 9216 + 1024) + (unsigned int)(tir * 4); + #pragma unroll + for (int p = 0; p < 3; p++) { + int rowq_ld = ((row_begin + p * rstep + rg) < (last_row) ? (row_begin + p * rstep + rg) : (last_row)); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.cg.shared::cta.global [%1], [%2], 16;\n\t" + "}" + :: "r"((p < n_iters) ? 1 : 0), "r"(sg + (unsigned int)(p * 1152)), "l"(g + ((unsigned long long)rowq_ld * (unsigned long long)g_stride + (unsigned long long)(tir * 8)))); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.cg.shared::cta.global [%1], [%2], 16;\n\t" + "}" + :: "r"((p < n_iters) ? 1 : 0), "r"(sg + (unsigned int)(p * 1152 + 512)), "l"(x + ((unsigned long long)rowq_ld * (unsigned long long)x_stride + (unsigned long long)(tir * 8)))); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.ca.shared::cta.global [%1], [%2], 4;\n\t" + "}" + :: "r"((p < n_iters) ? 1 : 0), "r"(sr + (unsigned int)(p * 1152)), "l"(r + rowq_ld)); + asm volatile("cp.async.commit_group;"); + } + int rowx = row_begin + 3 * rstep + rg; + int rowc = row_begin + rg; + int nok = (row_end - row_begin - rg + rstep - 1) / rstep; + if (nok > n_iters) { + nok = n_iters; + } + #pragma unroll 1 + for (int it = 0; it < n_iters; it += 2) { + unsigned int slx = (unsigned int)(it + 3 & 7) * 1152; + int rld = ((rowx) < (last_row) ? (rowx) : (last_row)); + unsigned long long gld = (unsigned long long)rld * (unsigned long long)g_stride + (unsigned long long)(tir * 8); + unsigned long long xld = (unsigned long long)rld * (unsigned long long)x_stride + (unsigned long long)(tir * 8); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.cg.shared::cta.global [%1], [%2], 16;\n\t" + "}" + :: "r"((it + 3 < n_iters) ? 1 : 0), "r"(sg + slx), "l"(g + gld)); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.cg.shared::cta.global [%1], [%2], 16;\n\t" + "}" + :: "r"((it + 3 < n_iters) ? 1 : 0), "r"(sg + slx + 512), "l"(x + xld)); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.ca.shared::cta.global [%1], [%2], 4;\n\t" + "}" + :: "r"((it + 3 < n_iters) ? 1 : 0), "r"(sr + slx), "l"(r + rld)); + asm volatile("cp.async.commit_group;"); + rowx = rowx + rstep; + unsigned int slx_0 = (unsigned int)(it + 3 + 1 & 7) * 1152; + int rld_0 = ((rowx) < (last_row) ? (rowx) : (last_row)); + unsigned long long gld_0 = (unsigned long long)rld_0 * (unsigned long long)g_stride + (unsigned long long)(tir * 8); + unsigned long long xld_0 = (unsigned long long)rld_0 * (unsigned long long)x_stride + (unsigned long long)(tir * 8); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.cg.shared::cta.global [%1], [%2], 16;\n\t" + "}" + :: "r"((it + 3 + 1 < n_iters) ? 1 : 0), "r"(sg + slx_0), "l"(g + gld_0)); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.cg.shared::cta.global [%1], [%2], 16;\n\t" + "}" + :: "r"((it + 3 + 1 < n_iters) ? 1 : 0), "r"(sg + slx_0 + 512), "l"(x + xld_0)); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.ca.shared::cta.global [%1], [%2], 4;\n\t" + "}" + :: "r"((it + 3 + 1 < n_iters) ? 1 : 0), "r"(sr + slx_0), "l"(r + rld_0)); + asm volatile("cp.async.commit_group;"); + rowx = rowx + rstep; + asm volatile("cp.async.wait_group 3;"); + unsigned int slc = (unsigned int)(it & 7) * 1152; + unsigned int slb = (unsigned int)(it + 1 & 7) * 1152; + __int128_t _smem_raw_3; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_3) : "r"(sg + slc)); + gwA[0] = reinterpret_cast(&_smem_raw_3)[0]; + gwA[1] = reinterpret_cast(&_smem_raw_3)[1]; + gwA[2] = reinterpret_cast(&_smem_raw_3)[2]; + gwA[3] = reinterpret_cast(&_smem_raw_3)[3]; + __int128_t _smem_raw_4; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_4) : "r"(sg + slc + 512)); + xwA[0] = reinterpret_cast(&_smem_raw_4)[0]; + xwA[1] = reinterpret_cast(&_smem_raw_4)[1]; + xwA[2] = reinterpret_cast(&_smem_raw_4)[2]; + xwA[3] = reinterpret_cast(&_smem_raw_4)[3]; + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&rrA[0])) : "r"(sr + slc)); + __int128_t _smem_raw_5; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_5) : "r"(sg + slb)); + gwB[0] = reinterpret_cast(&_smem_raw_5)[0]; + gwB[1] = reinterpret_cast(&_smem_raw_5)[1]; + gwB[2] = reinterpret_cast(&_smem_raw_5)[2]; + gwB[3] = reinterpret_cast(&_smem_raw_5)[3]; + __int128_t _smem_raw_6; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_6) : "r"(sg + slb + 512)); + xwB[0] = reinterpret_cast(&_smem_raw_6)[0]; + xwB[1] = reinterpret_cast(&_smem_raw_6)[1]; + xwB[2] = reinterpret_cast(&_smem_raw_6)[2]; + xwB[3] = reinterpret_cast(&_smem_raw_6)[3]; + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&rrB[0])) : "r"(sr + slb)); + if (it + 2 <= nok) { + float dq = 0.0f; + float gwr[8]; + #pragma unroll + for (int v_1 = 0; v_1 < 1; v_1++) { + unsigned int wvd[4]; + #pragma unroll + for (int e = 0; e < 4; e++) { + unsigned int gwd = gwA[v_1 * 4 + e]; + unsigned int xwd = xwA[v_1 * 4 + e]; + float g0 = __uint_as_float(gwd << 16); + float g1 = __uint_as_float(gwd & 4294901760u); + float x0 = __uint_as_float(xwd << 16); + float x1 = __uint_as_float(xwd & 4294901760u); + float2 wq = make_float2(__uint_as_float(ww[v_1 * 4 + e] << 16), __uint_as_float(ww[v_1 * 4 + e] & 4294901760u)); + float2 _f32x2_mul_a_7 = make_float2(g0, g1); + float2 gwq; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&gwq) : "l"(*(const unsigned long long*)&_f32x2_mul_a_7), "l"(*(const unsigned long long*)&wq)); + gwr[v_1 * 8 + 2 * e] = gwq.x; + gwr[v_1 * 8 + 2 * e + 1] = gwq.y; + dq = __fmaf_rn(gwq.x, x0, dq); + dq = __fmaf_rn(gwq.y, x1, dq); + } + } + float dots[1]; + float _warp_reduce_8 = dq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_8 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_8, offset); + dots[0] = _warp_reduce_8; + float rq = rrA[0]; + float cq = dots[0] * 0.00390625f; + float kx = rq * rq * rq * cq; + float nkx = -kx; + unsigned long long dx_base = (unsigned long long)rowc * 256; + float2 rr = make_float2(rq, rq); + float2 rrn = make_float2(-rq, -rq); + float2 nn = make_float2(nkx, nkx); + #pragma unroll + for (int v_2 = 0; v_2 < 1; v_2++) { + float out[8]; + unsigned int wv[4]; + #pragma unroll + for (int e_1 = 0; e_1 < 4; e_1++) { + unsigned int gwd_1 = gwA[v_2 * 4 + e_1]; + unsigned int xwd_1 = xwA[v_2 * 4 + e_1]; + float g0_1 = __uint_as_float(gwd_1 << 16); + float g1_1 = __uint_as_float(gwd_1 & 4294901760u); + float x0_1 = __uint_as_float(xwd_1 << 16); + float x1_1 = __uint_as_float(xwd_1 & 4294901760u); + float2 _f32x2_mul_a_9 = make_float2(g0_1, g1_1); + float2 _f32x2_mul_b_10 = make_float2(x0_1, x1_1); + float2 gx; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&gx) : "l"(*(const unsigned long long*)&_f32x2_mul_a_9), "l"(*(const unsigned long long*)&_f32x2_mul_b_10)); + float2 pn; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&pn) : "l"(*(const unsigned long long*)&gx), "l"(*(const unsigned long long*)&rrn)); + float2 pe = fma_f32x2_rn_noftz(gx, rr, pn); + float2 a2 = make_float2(acc[v_2 * 8 + 2 * e_1], acc[v_2 * 8 + 2 * e_1 + 1]); + float2 s2 = sub_f32x2_noftz(a2, pn); + float2 z2 = sub_f32x2_noftz(s2, a2); + float2 t2 = sub_f32x2_noftz(pe, add_f32x2_noftz(pn, z2)); + float2 c2 = add_f32x2_noftz(make_float2(comp[v_2 * 8 + 2 * e_1], comp[v_2 * 8 + 2 * e_1 + 1]), t2); + acc[v_2 * 8 + 2 * e_1] = s2.x; + acc[v_2 * 8 + 2 * e_1 + 1] = s2.y; + comp[v_2 * 8 + 2 * e_1] = c2.x; + comp[v_2 * 8 + 2 * e_1 + 1] = c2.y; + float2 _f32x2_mul_b_11 = make_float2(gwr[v_2 * 8 + 2 * e_1], gwr[v_2 * 8 + 2 * e_1 + 1]); + float2 q2; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&q2) : "l"(*(const unsigned long long*)&rr), "l"(*(const unsigned long long*)&_f32x2_mul_b_11)); + float2 o2 = fma_f32x2_rn_noftz(nn, make_float2(x0_1, x1_1), q2); + out[2 * e_1] = o2.x; + out[2 * e_1 + 1] = o2.y; + } + int col = (v_2 * 32 + tir) * 8; + unsigned int packed[4]; + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out[_lp*2 + 0], out[_lp*2+1 + 0])); + packed[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out[_lp*2 + 2], out[_lp*2+1 + 2])); + packed[_lp + 1] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out[_lp*2 + 4], out[_lp*2+1 + 4])); + packed[_lp + 2] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out[_lp*2 + 6], out[_lp*2+1 + 6])); + packed[_lp + 3] = *(uint32_t*)&_bf2; + } + reinterpret_cast(dx + (dx_base + (unsigned long long)col))[0] = reinterpret_cast(packed)[0]; + } + float dq_1 = 0.0f; + float gwr_0[8]; + #pragma unroll + for (int v_3 = 0; v_3 < 1; v_3++) { + unsigned int wvd_1[4]; + #pragma unroll + for (int e_2 = 0; e_2 < 4; e_2++) { + unsigned int gwd_2 = gwB[v_3 * 4 + e_2]; + unsigned int xwd_2 = xwB[v_3 * 4 + e_2]; + float g0_2 = __uint_as_float(gwd_2 << 16); + float g1_2 = __uint_as_float(gwd_2 & 4294901760u); + float x0_2 = __uint_as_float(xwd_2 << 16); + float x1_2 = __uint_as_float(xwd_2 & 4294901760u); + float2 wq_1 = make_float2(__uint_as_float(ww[v_3 * 4 + e_2] << 16), __uint_as_float(ww[v_3 * 4 + e_2] & 4294901760u)); + float2 _f32x2_mul_a_12 = make_float2(g0_2, g1_2); + float2 gwq_1; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&gwq_1) : "l"(*(const unsigned long long*)&_f32x2_mul_a_12), "l"(*(const unsigned long long*)&wq_1)); + gwr_0[v_3 * 8 + 2 * e_2] = gwq_1.x; + gwr_0[v_3 * 8 + 2 * e_2 + 1] = gwq_1.y; + dq_1 = __fmaf_rn(gwq_1.x, x0_2, dq_1); + dq_1 = __fmaf_rn(gwq_1.y, x1_2, dq_1); + } + } + float dots_0[1]; + float _warp_reduce_13 = dq_1; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_13 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_13, offset); + dots_0[0] = _warp_reduce_13; + float rq_0 = rrB[0]; + float cq_0 = dots_0[0] * 0.00390625f; + float kx_0 = rq_0 * rq_0 * rq_0 * cq_0; + float nkx_0 = -kx_0; + unsigned long long dx_base_0 = (unsigned long long)(rowc + rstep) * 256; + float2 rr_0 = make_float2(rq_0, rq_0); + float2 rrn_0 = make_float2(-rq_0, -rq_0); + float2 nn_0 = make_float2(nkx_0, nkx_0); + #pragma unroll + for (int v_4 = 0; v_4 < 1; v_4++) { + float out_1[8]; + unsigned int wv_1[4]; + #pragma unroll + for (int e_3 = 0; e_3 < 4; e_3++) { + unsigned int gwd_3 = gwB[v_4 * 4 + e_3]; + unsigned int xwd_3 = xwB[v_4 * 4 + e_3]; + float g0_3 = __uint_as_float(gwd_3 << 16); + float g1_3 = __uint_as_float(gwd_3 & 4294901760u); + float x0_3 = __uint_as_float(xwd_3 << 16); + float x1_3 = __uint_as_float(xwd_3 & 4294901760u); + float2 _f32x2_mul_a_14 = make_float2(g0_3, g1_3); + float2 _f32x2_mul_b_15 = make_float2(x0_3, x1_3); + float2 gx_1; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&gx_1) : "l"(*(const unsigned long long*)&_f32x2_mul_a_14), "l"(*(const unsigned long long*)&_f32x2_mul_b_15)); + float2 pn_1; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&pn_1) : "l"(*(const unsigned long long*)&gx_1), "l"(*(const unsigned long long*)&rrn_0)); + float2 pe_1 = fma_f32x2_rn_noftz(gx_1, rr_0, pn_1); + float2 a2_1 = make_float2(acc[v_4 * 8 + 2 * e_3], acc[v_4 * 8 + 2 * e_3 + 1]); + float2 s2_1 = sub_f32x2_noftz(a2_1, pn_1); + float2 z2_1 = sub_f32x2_noftz(s2_1, a2_1); + float2 t2_1 = sub_f32x2_noftz(pe_1, add_f32x2_noftz(pn_1, z2_1)); + float2 c2_1 = add_f32x2_noftz(make_float2(comp[v_4 * 8 + 2 * e_3], comp[v_4 * 8 + 2 * e_3 + 1]), t2_1); + acc[v_4 * 8 + 2 * e_3] = s2_1.x; + acc[v_4 * 8 + 2 * e_3 + 1] = s2_1.y; + comp[v_4 * 8 + 2 * e_3] = c2_1.x; + comp[v_4 * 8 + 2 * e_3 + 1] = c2_1.y; + float2 _f32x2_mul_b_16 = make_float2(gwr_0[v_4 * 8 + 2 * e_3], gwr_0[v_4 * 8 + 2 * e_3 + 1]); + float2 q2_1; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&q2_1) : "l"(*(const unsigned long long*)&rr_0), "l"(*(const unsigned long long*)&_f32x2_mul_b_16)); + float2 o2_1 = fma_f32x2_rn_noftz(nn_0, make_float2(x0_3, x1_3), q2_1); + out_1[2 * e_3] = o2_1.x; + out_1[2 * e_3 + 1] = o2_1.y; + } + int col_1 = (v_4 * 32 + tir) * 8; + unsigned int packed_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_1[_lp*2 + 0], out_1[_lp*2+1 + 0])); + packed_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_1[_lp*2 + 2], out_1[_lp*2+1 + 2])); + packed_1[_lp + 1] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_1[_lp*2 + 4], out_1[_lp*2+1 + 4])); + packed_1[_lp + 2] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_1[_lp*2 + 6], out_1[_lp*2+1 + 6])); + packed_1[_lp + 3] = *(uint32_t*)&_bf2; + } + reinterpret_cast(dx + (dx_base_0 + (unsigned long long)col_1))[0] = reinterpret_cast(packed_1)[0]; + } + } else if (it < nok) { + float dq_2 = 0.0f; + float gwr_1[8]; + #pragma unroll + for (int v_5 = 0; v_5 < 1; v_5++) { + unsigned int wvd_2[4]; + #pragma unroll + for (int e_4 = 0; e_4 < 4; e_4++) { + unsigned int gwd_4 = gwA[v_5 * 4 + e_4]; + unsigned int xwd_4 = xwA[v_5 * 4 + e_4]; + float g0_4 = __uint_as_float(gwd_4 << 16); + float g1_4 = __uint_as_float(gwd_4 & 4294901760u); + float x0_4 = __uint_as_float(xwd_4 << 16); + float x1_4 = __uint_as_float(xwd_4 & 4294901760u); + float2 wq_2 = make_float2(__uint_as_float(ww[v_5 * 4 + e_4] << 16), __uint_as_float(ww[v_5 * 4 + e_4] & 4294901760u)); + float2 _f32x2_mul_a_17 = make_float2(g0_4, g1_4); + float2 gwq_2; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&gwq_2) : "l"(*(const unsigned long long*)&_f32x2_mul_a_17), "l"(*(const unsigned long long*)&wq_2)); + gwr_1[v_5 * 8 + 2 * e_4] = gwq_2.x; + gwr_1[v_5 * 8 + 2 * e_4 + 1] = gwq_2.y; + dq_2 = __fmaf_rn(gwq_2.x, x0_4, dq_2); + dq_2 = __fmaf_rn(gwq_2.y, x1_4, dq_2); + } + } + float dots_1[1]; + float _warp_reduce_18 = dq_2; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_18 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_18, offset); + dots_1[0] = _warp_reduce_18; + float rq_1 = rrA[0]; + float cq_1 = dots_1[0] * 0.00390625f; + float kx_1 = rq_1 * rq_1 * rq_1 * cq_1; + float nkx_1 = -kx_1; + unsigned long long dx_base_1 = (unsigned long long)rowc * 256; + float2 rr_1 = make_float2(rq_1, rq_1); + float2 rrn_1 = make_float2(-rq_1, -rq_1); + float2 nn_1 = make_float2(nkx_1, nkx_1); + #pragma unroll + for (int v_6 = 0; v_6 < 1; v_6++) { + float out_2[8]; + unsigned int wv_2[4]; + #pragma unroll + for (int e_5 = 0; e_5 < 4; e_5++) { + unsigned int gwd_5 = gwA[v_6 * 4 + e_5]; + unsigned int xwd_5 = xwA[v_6 * 4 + e_5]; + float g0_5 = __uint_as_float(gwd_5 << 16); + float g1_5 = __uint_as_float(gwd_5 & 4294901760u); + float x0_5 = __uint_as_float(xwd_5 << 16); + float x1_5 = __uint_as_float(xwd_5 & 4294901760u); + float2 _f32x2_mul_a_19 = make_float2(g0_5, g1_5); + float2 _f32x2_mul_b_20 = make_float2(x0_5, x1_5); + float2 gx_2; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&gx_2) : "l"(*(const unsigned long long*)&_f32x2_mul_a_19), "l"(*(const unsigned long long*)&_f32x2_mul_b_20)); + float2 pn_2; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&pn_2) : "l"(*(const unsigned long long*)&gx_2), "l"(*(const unsigned long long*)&rrn_1)); + float2 pe_2 = fma_f32x2_rn_noftz(gx_2, rr_1, pn_2); + float2 a2_2 = make_float2(acc[v_6 * 8 + 2 * e_5], acc[v_6 * 8 + 2 * e_5 + 1]); + float2 s2_2 = sub_f32x2_noftz(a2_2, pn_2); + float2 z2_2 = sub_f32x2_noftz(s2_2, a2_2); + float2 t2_2 = sub_f32x2_noftz(pe_2, add_f32x2_noftz(pn_2, z2_2)); + float2 c2_2 = add_f32x2_noftz(make_float2(comp[v_6 * 8 + 2 * e_5], comp[v_6 * 8 + 2 * e_5 + 1]), t2_2); + acc[v_6 * 8 + 2 * e_5] = s2_2.x; + acc[v_6 * 8 + 2 * e_5 + 1] = s2_2.y; + comp[v_6 * 8 + 2 * e_5] = c2_2.x; + comp[v_6 * 8 + 2 * e_5 + 1] = c2_2.y; + float2 _f32x2_mul_b_21 = make_float2(gwr_1[v_6 * 8 + 2 * e_5], gwr_1[v_6 * 8 + 2 * e_5 + 1]); + float2 q2_2; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&q2_2) : "l"(*(const unsigned long long*)&rr_1), "l"(*(const unsigned long long*)&_f32x2_mul_b_21)); + float2 o2_2 = fma_f32x2_rn_noftz(nn_1, make_float2(x0_5, x1_5), q2_2); + out_2[2 * e_5] = o2_2.x; + out_2[2 * e_5 + 1] = o2_2.y; + } + int col_2 = (v_6 * 32 + tir) * 8; + unsigned int packed_2[4]; + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_2[_lp*2 + 0], out_2[_lp*2+1 + 0])); + packed_2[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_2[_lp*2 + 2], out_2[_lp*2+1 + 2])); + packed_2[_lp + 1] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_2[_lp*2 + 4], out_2[_lp*2+1 + 4])); + packed_2[_lp + 2] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_2[_lp*2 + 6], out_2[_lp*2+1 + 6])); + packed_2[_lp + 3] = *(uint32_t*)&_bf2; + } + reinterpret_cast(dx + (dx_base_1 + (unsigned long long)col_2))[0] = reinterpret_cast(packed_2)[0]; + } + } + rowc = rowc + 2 * rstep; + } + asm volatile("cp.async.wait_group 0;"); + } + unsigned long long pbase = (unsigned long long)c * 256; + #pragma unroll + for (int v_7 = 0; v_7 < 1; v_7++) { + int fcol = (v_7 * 32 + tir) * 8; + #pragma unroll + for (int e_6 = 0; e_6 < 8; e_6++) { + row_fold[rg * 256 + fcol + e_6] = acc[v_7 * 8 + e_6]; + } + } + __syncthreads(); + float fold_acc[2]; + #pragma unroll + for (int e_7 = 0; e_7 < 2; e_7++) { + fold_acc[e_7] = 0.0f; + } + float fold_c[2]; + #pragma unroll + for (int e_8 = 0; e_8 < 2; e_8++) { + fold_c[e_8] = 0.0f; + } + #pragma unroll + for (int k_1 = 0; k_1 < 4; k_1++) { + #pragma unroll + for (int e_9 = 0; e_9 < 2; e_9++) { + float d2 = row_fold[k_1 * 256 + tid * 2 + e_9]; + float fs = fold_acc[e_9] + d2; + float fb = fs - fold_acc[e_9]; + fold_c[e_9] = fold_c[e_9] + (fold_acc[e_9] - (fs - fb) + (d2 - fb)); + fold_acc[e_9] = fs; + } + } + __syncthreads(); + #pragma unroll + for (int e5 = 0; e5 < 8; e5++) { + row_fold[rg * 256 + tir * 8 + e5] = comp[e5]; + } + __syncthreads(); + #pragma unroll + for (int k5 = 0; k5 < 4; k5++) { + #pragma unroll + for (int e_10 = 0; e_10 < 2; e_10++) { + float d5 = row_fold[k5 * 256 + tid * 2 + e_10]; + fold_c[e_10] = fold_c[e_10] + d5; + } + } + #pragma unroll + for (int e_11 = 0; e_11 < 2; e_11++) { + fold_acc[e_11] = fold_acc[e_11] + fold_c[e_11]; + } + int pcol = tid * 2; + { + float2 _v2 = make_float2(fold_acc[0 + 0], fold_acc[0 + 1]); + *reinterpret_cast(partial + pbase + (unsigned long long)pcol) = _v2; + } + if (tid == 0) { + flags[claim_it & 1] = (int)nxt; + } + __syncthreads(); + c = flags[claim_it & 1]; + } + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + if (tid == 0) { + if (c == n_chunks + (int)gridDim.x - 1) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } +} + +} // extern "C" diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_binding.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_binding.cu new file mode 100644 index 00000000000..d4496992831 --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_binding.cu @@ -0,0 +1,146 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_det_train_0f0b782f1e538caad3ce, for the device source that defines it. Do not edit. +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_rmsnorm_det_train_0f0b782f1e538caad3ce(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, float* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0); + +namespace cake_binding_kernel_cake_rmsnorm_det_train_0f0b782f1e538caad3ce { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_det_train_0f0b782f1e538caad3ce"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckRank(const TensorView& t, const char* name, int rank) { + TVM_FFI_CHECK(t.ndim() >= rank, ValueError) + << KERNEL_NAME << ": '" << name << "' needs at least " << rank << " dims, got " << t.ndim(); +} + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +void Run(TensorView a_g, TensorView a_x, TensorView a_w, TensorView a_r, TensorView a_dx, TensorView a_dw, TensorView a_partial, TensorView a_counters, TensorView a_g_h, int64_t a_T, int64_t a_n_chunks, int64_t a_rows_per_chunk, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_g.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_g, "g"); + check_dtype(a_g, DLDataType{kDLBfloat, 16, 1}, "g"); + check_cuda_tensor(a_x, "x"); + check_dtype(a_x, DLDataType{kDLBfloat, 16, 1}, "x"); + check_cuda_tensor(a_w, "w"); + check_dtype(a_w, DLDataType{kDLFloat, 32, 1}, "w"); + check_cuda_tensor(a_r, "r"); + check_dtype(a_r, DLDataType{kDLFloat, 32, 1}, "r"); + check_cuda_tensor(a_dx, "dx"); + check_dtype(a_dx, DLDataType{kDLBfloat, 16, 1}, "dx"); + check_cuda_tensor(a_dw, "dw"); + check_dtype(a_dw, DLDataType{kDLFloat, 32, 1}, "dw"); + check_cuda_tensor(a_partial, "partial"); + check_dtype(a_partial, DLDataType{kDLFloat, 32, 1}, "partial"); + check_cuda_tensor(a_counters, "counters"); + check_dtype(a_counters, DLDataType{kDLUInt, 32, 1}, "counters"); + check_cuda_tensor(a_g_h, "g_h"); + check_dtype(a_g_h, DLDataType{kDLBfloat, 16, 1}, "g_h"); + TVM_FFI_CHECK(a_T >= -2147483648LL && a_T <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'T' = " << a_T + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_n_chunks >= -2147483648LL && a_n_chunks <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'n_chunks' = " << a_n_chunks + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_rows_per_chunk >= -2147483648LL && a_rows_per_chunk <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'rows_per_chunk' = " << a_rows_per_chunk + << " is outside the i32 range [-2147483648, 2147483647]"; + CheckRank(a_g, "g", 1); + CheckUnitInnerStride(a_g, "g"); + CheckRank(a_x, "x", 1); + CheckUnitInnerStride(a_x, "x"); + check_same_device(a_x, a_g, "x", "g"); + check_contiguous(a_w, "w"); + check_same_device(a_w, a_g, "w", "g"); + check_contiguous(a_r, "r"); + check_same_device(a_r, a_g, "r", "g"); + check_contiguous(a_dx, "dx"); + check_same_device(a_dx, a_g, "dx", "g"); + check_contiguous(a_dw, "dw"); + check_same_device(a_dw, a_g, "dw", "g"); + check_contiguous(a_partial, "partial"); + check_same_device(a_partial, a_g, "partial", "g"); + check_contiguous(a_counters, "counters"); + check_same_device(a_counters, a_g, "counters", "g"); + CheckRank(a_g_h, "g_h", 1); + CheckUnitInnerStride(a_g_h, "g_h"); + check_same_device(a_g_h, a_g, "g_h", "g"); + TVM_FFI_CHECK((a_n_chunks >= 1LL), ValueError) + << KERNEL_NAME << ": requires (n_chunks >= 1)"; + TVM_FFI_CHECK((a_rows_per_chunk >= 1LL), ValueError) + << KERNEL_NAME << ": requires (rows_per_chunk >= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_rmsnorm_det_train_0f0b782f1e538caad3ce), cudaFuncAttributeMaxDynamicSharedMemorySize, 99840), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + void* v_g = DataPtr(a_g); + void* v_x = DataPtr(a_x); + void* v_w = DataPtr(a_w); + void* v_r = DataPtr(a_r); + void* v_dx = DataPtr(a_dx); + void* v_dw = DataPtr(a_dw); + void* v_partial = DataPtr(a_partial); + void* v_counters = DataPtr(a_counters); + void* v_g_h = DataPtr(a_g_h); + int32_t v_T = static_cast(a_T); + int32_t v_n_chunks = static_cast(a_n_chunks); + int32_t v_rows_per_chunk = static_cast(a_rows_per_chunk); + int64_t v_g_stride_p0 = a_g.stride(0); + int64_t v_x_stride_p0 = a_x.stride(0); + int64_t v_g_h_stride_p0 = a_g_h.stride(0); + void* kernel_args[] = {&v_g, &v_x, &v_w, &v_r, &v_dx, &v_dw, &v_partial, &v_counters, &v_g_h, &v_T, &v_n_chunks, &v_rows_per_chunk, &v_g_stride_p0, &v_x_stride_p0, &v_g_h_stride_p0}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(384, 1, 1); + config.dynamicSmemBytes = 99840; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_rmsnorm_det_train_0f0b782f1e538caad3ce), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_rmsnorm_det_train_0f0b782f1e538caad3ce + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_det_train_0f0b782f1e538caad3ce::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_kernel.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_kernel.cu new file mode 100644 index 00000000000..3b6fe8c1e10 --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_0f0b782f1e538caad3ce_kernel.cu @@ -0,0 +1,832 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define SMEM_STAGE_BUF_OFF 1024 +#define SMEM_STAGE_BUF_STAGE_BYTES 73728 +#define SMEM_STAGE_BUF_STRIDE 73728 +#define SMEM_ROW_FOLD_OFF 1024 +#define SMEM_ROW_FOLD_STAGE_BYTES 24576 +#define SMEM_ROW_FOLD_STRIDE 24576 +#define SMEM_RED_FOLD_OFF 25600 +#define SMEM_RED_FOLD_STAGE_BYTES 6144 +#define SMEM_RED_FOLD_STRIDE 6144 +#define SMEM_DOT_SLOTS_OFF 74752 +#define SMEM_DOT_SLOTS_STAGE_BYTES 96 +#define SMEM_DOT_SLOTS_STRIDE 96 +#define SMEM_PART_SM_OFF 74880 +#define SMEM_PART_SM_STAGE_BYTES 16 +#define SMEM_PART_SM_STRIDE 16 +#define SMEM_DESC_OFF 99584 +#define SMEM_DESC_STAGE_BYTES 60 +#define SMEM_DESC_STRIDE 60 +#define SMEM_RTAB_OFF 99712 +#define SMEM_RTAB_STAGE_BYTES 128 +#define SMEM_RTAB_STRIDE 128 +#define SMEM_FOLD_SM_OFF 75008 +#define SMEM_FOLD_SM_STAGE_BYTES 24576 +#define SMEM_FOLD_SM_STRIDE 24576 +#define SMEM_TOTAL 99840 +#define THREADS 384 +#define LAUNCH_MIN_BLOCKS 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_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"); +} + + + +extern "C" { + +__global__ __launch_bounds__(384, LAUNCH_MIN_BLOCKS) void +kernel_cake_rmsnorm_det_train_0f0b782f1e538caad3ce(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, float* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define bars_addr (mbar_base + 0) + #define dbar_addr (mbar_base + 96) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + uint8_t* stage_buf = reinterpret_cast(smem_raw + 1024); + const int stage_buf_addr = smem + 1024; + float* row_fold = reinterpret_cast(smem_raw + 1024); + const int row_fold_addr = smem + 1024; + float* red_fold = reinterpret_cast(smem_raw + 25600); + const int red_fold_addr = smem + 25600; + float* dot_slots = reinterpret_cast(smem_raw + 74752); + const int dot_slots_addr = smem + 74752; + float* part_sm = reinterpret_cast(smem_raw + 74880); + const int part_sm_addr = smem + 74880; + int* desc = reinterpret_cast(smem_raw + 99584); + const int desc_addr = smem + 99584; + int* rtab = reinterpret_cast(smem_raw + 99712); + const int rtab_addr = smem + 99712; + float* fold_sm = reinterpret_cast(smem_raw + 75008); + const int fold_sm_addr = smem + 75008; + + // Mbarrier init (2 pipeline groups, 0 ordered-sequence groups, 17 barriers) + // Mbarriers at smem_raw[0..136) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // bars: 12 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + mbarrier_init(smem + 8, 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); + // dbar: 5 barriers, init_count=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); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncthreads(); + + // === Task calls (declaration order) === + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + long long g_stride = g_stride_p0; + long long x_stride = x_stride_p0; + long long gh_stride = g_h_stride_p0; + int c = bid; + int rg = tid / 96; + int tir = tid - rg * 96; + unsigned int claim[1]; + if (tid == 0) { + unsigned int _atomic_fetch_add_0 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_0; + int S = (T + 3) / 4; + int _min_1 = ((rows_per_chunk) < (S) ? (rows_per_chunk) : (S)); + int sl = ((_min_1) > (1) ? (_min_1) : (1)); + int share = (S + sl - 1) / sl; + int tail = (((share + 11) / 12) > (1) ? ((share + 11) / 12) : (1)); + int size = (share + 1) / 2; + int rem = S; + int base = 0; + int first = 0; + int nr = 0; + int walk = 1; + while (walk == 1) { + walk = 0; + if (nr < 5) { + if (size > tail) { + if (rem >= sl * size) { + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = size; + first += sl; + base += sl * size; + rem -= sl * size; + size = (size + 1) / 2; + nr += 1; + walk = 1; + } + } + } + } + int q = rem / sl; + int m = rem - q * sl; + if (m > 0) { + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = q + 1; + first += m; + base += m * (q + 1); + nr += 1; + } + if (q > 0) { + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = q; + first += sl - m; + base += (sl - m) * q; + nr += 1; + } + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = 0; + rtab[31] = nr + 1; + } + int ring = (int)stage_buf_addr; + int my_g = ring + rg * 3072 + tir * 16; + float wf[16]; + #pragma unroll + for (int v = 0; v < 2; v++) { + int wcol = (v * 96 + tir) * 8; + { + unsigned _ldv8_2_0; + unsigned _ldv8_2_1; + unsigned _ldv8_2_2; + unsigned _ldv8_2_3; + unsigned _ldv8_2_4; + unsigned _ldv8_2_5; + unsigned _ldv8_2_6; + unsigned _ldv8_2_7; + asm volatile( + "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_2_0), "=r"(_ldv8_2_1), "=r"(_ldv8_2_2), "=r"(_ldv8_2_3), "=r"(_ldv8_2_4), "=r"(_ldv8_2_5), "=r"(_ldv8_2_6), "=r"(_ldv8_2_7) : "l"((const void*)(w + ((unsigned long long)wcol))) : "memory"); + wf[v * 8 + 0] = __uint_as_float(_ldv8_2_0); + wf[v * 8 + 1] = __uint_as_float(_ldv8_2_1); + wf[v * 8 + 2] = __uint_as_float(_ldv8_2_2); + wf[v * 8 + 3] = __uint_as_float(_ldv8_2_3); + wf[v * 8 + 4] = __uint_as_float(_ldv8_2_4); + wf[v * 8 + 5] = __uint_as_float(_ldv8_2_5); + wf[v * 8 + 6] = __uint_as_float(_ldv8_2_6); + wf[v * 8 + 7] = __uint_as_float(_ldv8_2_7); + } + } + int rng[2]; + int nd[3]; + if (tid == 0) { + int cid0 = (int)claim[0]; + int nr_1 = rtab[31]; + int r_0 = 0; + int walk_1 = 1; + while (walk_1 == 1) { + walk_1 = 0; + if (nr_1 > r_0 + 2) { + if (cid0 >= rtab[(r_0 + 1) * 3]) { + r_0 += 1; + walk_1 = 1; + } + } + } + if (cid0 >= rtab[(nr_1 - 1) * 3]) { + rng[0] = T; + rng[1] = T; + } else { + int first_1 = rtab[r_0 * 3]; + int base_1 = rtab[r_0 * 3 + 1]; + int stages = rtab[r_0 * 3 + 2]; + rng[0] = (base_1 + (cid0 - first_1) * stages) * 4; + int _min_3 = ((rng[0] + stages * 4) < (T) ? (rng[0] + stages * 4) : (T)); + rng[1] = _min_3; + } + int entry = 0; + int word = entry * 3; + desc[word] = cid0; + desc[word + 1] = rng[0]; + desc[word + 2] = rng[1]; + mbarrier_arrive(dbar_addr + (entry) * 8); + if (rng[0] < T) { + unsigned int _atomic_fetch_add_4 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_4; + } else { + int nr_0 = rtab[31]; + int last = rtab[(nr_0 - 1) * 3] + num_bids - 1; + if (T == 0) { + last = 0; + } + if (cid0 == last) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } + } + int entry_1 = 0; + mbarrier_wait(dbar_addr + (entry_1) * 8, 0); + int word_1 = entry_1 * 3; + int w0 = desc[word_1]; + int w1 = desc[word_1 + 1]; + int w2 = desc[word_1 + 2]; + int _warp_redux_5; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_5) : "r"(w0)); + nd[0] = _warp_redux_5; + int _warp_redux_6; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_6) : "r"(w1)); + nd[1] = _warp_redux_6; + int _warp_redux_7; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_7) : "r"(w2)); + nd[2] = _warp_redux_7; + int c0 = nd[0]; + int cb = nd[1]; + int ce = nd[2]; + int ps[3]; + ps[0] = nd[1]; + ps[1] = nd[2]; + int desc_n = 0; + #pragma unroll + for (int s = 0; s < 3; s++) { + ps[2] = 0; + if (ps[0] < ps[1]) { + int first_2 = ps[0]; + int rows_k = ((ps[1] - first_2) < (4) ? (ps[1] - first_2) : (4)); + int sbase = ring + s * 24576; + int gstage = rg + 4 * s; + if (tir == 0) { + if (rg < rows_k) { + mbarrier_arrive_expect_tx(bars_addr + (gstage) * 8, 6144); + int row_g = first_2 + rg; + long long g_srcg = (long long)row_g * g_stride; + long long x_srcg = (long long)row_g * x_stride; + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase + rg * 3072), "l"(reinterpret_cast(g + g_srcg)), "r"((unsigned)(3072)), "r"(bars_addr + (gstage) * 8), + "l"(0x14F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase + (rg + 4) * 3072), "l"(reinterpret_cast(x + x_srcg)), "r"((unsigned)(3072)), "r"(bars_addr + (gstage) * 8), + "l"(0x14F0000000000000ULL) : "memory"); + } else { + mbarrier_arrive(bars_addr + (gstage) * 8); + } + } + ps[0] = ps[0] + 4; + ps[2] = (int)(ps[0] >= ps[1]); + if (ps[0] >= ps[1]) { + if (tid == 0) { + int cid = (int)claim[0]; + int nr_2 = rtab[31]; + int r_0_1 = 0; + int walk_2 = 1; + while (walk_2 == 1) { + walk_2 = 0; + if (nr_2 > r_0_1 + 2) { + if (cid >= rtab[(r_0_1 + 1) * 3]) { + r_0_1 += 1; + walk_2 = 1; + } + } + } + if (cid >= rtab[(nr_2 - 1) * 3]) { + rng[0] = T; + rng[1] = T; + } else { + int first_0 = rtab[r_0_1 * 3]; + int base_2 = rtab[r_0_1 * 3 + 1]; + int stages_1 = rtab[r_0_1 * 3 + 2]; + rng[0] = (base_2 + (cid - first_0) * stages_1) * 4; + int _min_8 = ((rng[0] + stages_1 * 4) < (T) ? (rng[0] + stages_1 * 4) : (T)); + rng[1] = _min_8; + } + int entry_1_1 = (desc_n + 1) % 5; + int word_2 = entry_1_1 * 3; + desc[word_2] = cid; + desc[word_2 + 1] = rng[0]; + desc[word_2 + 2] = rng[1]; + mbarrier_arrive(dbar_addr + (entry_1_1) * 8); + if (rng[0] < T) { + unsigned int _atomic_fetch_add_9 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_9; + } else { + int nr_0_1 = rtab[31]; + int last_1 = rtab[(nr_0_1 - 1) * 3] + num_bids - 1; + if (T == 0) { + last_1 = 0; + } + if (cid == last_1) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } + } + int entry_0 = (desc_n + 1) % 5; + int _parity_dbar = (int)((unsigned int)(desc_n + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_0) * 8, _parity_dbar); + int word_1_1 = entry_0 * 3; + int w0_2 = desc[word_1_1]; + int w1_3 = desc[word_1_1 + 1]; + int w2_4 = desc[word_1_1 + 2]; + int _warp_redux_10; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_10) : "r"(w0_2)); + nd[0] = _warp_redux_10; + int _warp_redux_11; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_11) : "r"(w1_3)); + nd[1] = _warp_redux_11; + int _warp_redux_12; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_12) : "r"(w2_4)); + nd[2] = _warp_redux_12; + ps[0] = nd[1]; + ps[1] = nd[2]; + } + } + int desc_step = ps[2]; + desc_n += desc_step; + } + float acc[16]; + #pragma unroll + for (int j = 0; j < 16; j++) { + acc[j] = 0.0f; + } + float cmp[16]; + #pragma unroll + for (int jc = 0; jc < 16; jc++) { + cmp[jc] = 0.0f; + } + float hi[16]; + #pragma unroll + for (int jh = 0; jh < 16; jh++) { + hi[jh] = 0.0f; + } + int n_cons = 0; + int g_step = 0; + float r_cur = 0.0f; + if (cb < ce) { + int row0 = ((cb + rg) < (ce - 1) ? (cb + rg) : (ce - 1)); + r_cur = r[row0]; + } + unsigned int gw[8]; + unsigned int xw[8]; + float out[16]; + if (cb >= ce) { + if (c0 < n_chunks) { + unsigned long long pbase0 = (unsigned long long)c0 * 1536; + float fold0[4]; + #pragma unroll + for (int e0 = 0; e0 < 4; e0++) { + fold0[e0] = 0.0f; + } + int pcol0 = tid * 4; + { + float4 _v4 = make_float4(fold0[0 + 0], fold0[0 + 1], fold0[0 + 2], fold0[0 + 3]); + *reinterpret_cast(partial + pbase0 + (unsigned long long)pcol0) = _v4; + } + } + } + while (cb < ce) { + int n_steps_c = (ce - cb + 3) / 4; + #pragma unroll 1 + for (int k = 0; k < n_steps_c; k++) { + int slot = g_step % 3; + int par = (g_step & 1) * 12; + int row = cb + k * 4 + rg; + int row_n = ((row + 4) < (ce - 1) ? (row + 4) : (ce - 1)); + if (n_steps_c <= k + 1) { + int entry_0_1 = (n_cons + 1) % 5; + int _parity_dbar_1 = (int)((unsigned int)(n_cons + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_0_1) * 8, _parity_dbar_1); + int word_0 = entry_0_1 * 3; + int w0_0 = desc[word_0]; + int w1_0 = desc[word_0 + 1]; + int w2_0 = desc[word_0 + 2]; + int _warp_redux_13; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_13) : "r"(w0_0)); + nd[0] = _warp_redux_13; + int _warp_redux_14; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_14) : "r"(w1_0)); + nd[1] = _warp_redux_14; + int _warp_redux_15; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_15) : "r"(w2_0)); + nd[2] = _warp_redux_15; + row_n = cb; + if (nd[1] < T) { + int _min_16 = ((nd[1] + rg) < (nd[2] - 1) ? (nd[1] + rg) : (nd[2] - 1)); + row_n = _min_16; + } + } + float r_nxt = r[row_n]; + int _stage_bars = rg + 4 * (int)((unsigned int)g_step % 3); + int _parity_bars = (int)((unsigned int)g_step / 3 % 2); + mbarrier_wait(bars_addr + (_stage_bars) * 8, _parity_bars); + int sg = my_g + slot * 24576; + #pragma unroll + for (int v_1 = 0; v_1 < 2; v_1++) { + __int128_t _smem_raw_17; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_17) : "r"(sg + v_1 * 1536)); + gw[v_1 * 4] = reinterpret_cast(&_smem_raw_17)[0]; + gw[(v_1 * 4) + 1] = reinterpret_cast(&_smem_raw_17)[1]; + gw[(v_1 * 4) + 2] = reinterpret_cast(&_smem_raw_17)[2]; + gw[(v_1 * 4) + 3] = reinterpret_cast(&_smem_raw_17)[3]; + __int128_t _smem_raw_18; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_18) : "r"(sg + (12288 + v_1 * 1536))); + xw[v_1 * 4] = reinterpret_cast(&_smem_raw_18)[0]; + xw[(v_1 * 4) + 1] = reinterpret_cast(&_smem_raw_18)[1]; + xw[(v_1 * 4) + 2] = reinterpret_cast(&_smem_raw_18)[2]; + xw[(v_1 * 4) + 3] = reinterpret_cast(&_smem_raw_18)[3]; + } + float dq = 0.0f; + #pragma unroll + for (int v_2 = 0; v_2 < 2; v_2++) { + #pragma unroll + for (int e = 0; e < 4; e++) { + unsigned int gwd = gw[v_2 * 4 + e]; + unsigned int xwd = xw[v_2 * 4 + e]; + float g0 = __uint_as_float(gwd << 16); + float g1 = __uint_as_float(gwd & 4294901760u); + float x0 = __uint_as_float(xwd << 16); + float x1 = __uint_as_float(xwd & 4294901760u); + dq = __fmaf_rn(g0 * wf[v_2 * 8 + 2 * e], x0, dq); + dq = __fmaf_rn(g1 * wf[v_2 * 8 + 2 * e + 1], x1, dq); + } + } + float dot = dq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + dot += __shfl_xor_sync(0xFFFFFFFF, dot, offset); + if (lane == 0) { + dot_slots[par + warp] = dot; + } + asm volatile("barrier.sync %0, 96;" :: "r"(1 + rg) : "memory"); + ps[2] = 0; + if (ps[0] < ps[1]) { + int first_3 = ps[0]; + int rows_k_1 = ((ps[1] - first_3) < (4) ? (ps[1] - first_3) : (4)); + int sbase_1 = ring + g_step % 3 * 24576; + int gstage_1 = rg + 4 * (g_step % 3); + if (tir == 0) { + if (rg < rows_k_1) { + mbarrier_arrive_expect_tx(bars_addr + (gstage_1) * 8, 6144); + int row_g_1 = first_3 + rg; + long long g_srcg_1 = (long long)row_g_1 * g_stride; + long long x_srcg_1 = (long long)row_g_1 * x_stride; + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase_1 + rg * 3072), "l"(reinterpret_cast(g + g_srcg_1)), "r"((unsigned)(3072)), "r"(bars_addr + (gstage_1) * 8), + "l"(0x14F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase_1 + (rg + 4) * 3072), "l"(reinterpret_cast(x + x_srcg_1)), "r"((unsigned)(3072)), "r"(bars_addr + (gstage_1) * 8), + "l"(0x14F0000000000000ULL) : "memory"); + } else { + mbarrier_arrive(bars_addr + (gstage_1) * 8); + } + } + ps[0] = ps[0] + 4; + ps[2] = (int)(ps[0] >= ps[1]); + if (ps[0] >= ps[1]) { + if (tid == 0) { + int cid_1 = (int)claim[0]; + int nr_3 = rtab[31]; + int r_0_2 = 0; + int walk_3 = 1; + while (walk_3 == 1) { + walk_3 = 0; + if (nr_3 > r_0_2 + 2) { + if (cid_1 >= rtab[(r_0_2 + 1) * 3]) { + r_0_2 += 1; + walk_3 = 1; + } + } + } + if (cid_1 >= rtab[(nr_3 - 1) * 3]) { + rng[0] = T; + rng[1] = T; + } else { + int first_1_1 = rtab[r_0_2 * 3]; + int base_3 = rtab[r_0_2 * 3 + 1]; + int stages_2 = rtab[r_0_2 * 3 + 2]; + rng[0] = (base_3 + (cid_1 - first_1_1) * stages_2) * 4; + int _min_19 = ((rng[0] + stages_2 * 4) < (T) ? (rng[0] + stages_2 * 4) : (T)); + rng[1] = _min_19; + } + int entry_1_2 = (desc_n + 1) % 5; + int word_2_1 = entry_1_2 * 3; + desc[word_2_1] = cid_1; + desc[word_2_1 + 1] = rng[0]; + desc[word_2_1 + 2] = rng[1]; + mbarrier_arrive(dbar_addr + (entry_1_2) * 8); + if (rng[0] < T) { + unsigned int _atomic_fetch_add_20 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_20; + } else { + int nr_0_2 = rtab[31]; + int last_2 = rtab[(nr_0_2 - 1) * 3] + num_bids - 1; + if (T == 0) { + last_2 = 0; + } + if (cid_1 == last_2) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } + } + int entry_0_2 = (desc_n + 1) % 5; + int _parity_dbar_2 = (int)((unsigned int)(desc_n + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_0_2) * 8, _parity_dbar_2); + int word_1_2 = entry_0_2 * 3; + int w0_2_1 = desc[word_1_2]; + int w1_3_1 = desc[word_1_2 + 1]; + int w2_4_1 = desc[word_1_2 + 2]; + int _warp_redux_21; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_21) : "r"(w0_2_1)); + nd[0] = _warp_redux_21; + int _warp_redux_22; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_22) : "r"(w1_3_1)); + nd[1] = _warp_redux_22; + int _warp_redux_23; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_23) : "r"(w2_4_1)); + nd[2] = _warp_redux_23; + ps[0] = nd[1]; + ps[1] = nd[2]; + } + } + int desc_step_1 = ps[2]; + desc_n += desc_step_1; + float tot = 0.0f; + #pragma unroll + for (int kk = 0; kk < 3; kk++) { + float d = dot_slots[par + rg * 3 + kk]; + tot += d; + } + float rq = r_cur; + float cq = __fdiv_rn(tot, 1536.0f); + float kx = rq * rq * rq * cq; + float nkx = -kx; + if (row < ce) { + unsigned long long dx_base = (unsigned long long)row * 1536; + #pragma unroll + for (int v_3 = 0; v_3 < 2; v_3++) { + #pragma unroll + for (int e_1 = 0; e_1 < 4; e_1++) { + unsigned int gwd2 = gw[v_3 * 4 + e_1]; + unsigned int xwd2 = xw[v_3 * 4 + e_1]; + float g2 = __uint_as_float(gwd2 << 16); + float g3 = __uint_as_float(gwd2 & 4294901760u); + float x2 = __uint_as_float(xwd2 << 16); + float x3 = __uint_as_float(xwd2 & 4294901760u); + acc[v_3 * 8 + 2 * e_1] = __fmaf_rn(g2 * x2, rq, acc[v_3 * 8 + 2 * e_1]); + acc[v_3 * 8 + 2 * e_1 + 1] = __fmaf_rn(g3 * x3, rq, acc[v_3 * 8 + 2 * e_1 + 1]); + out[v_3 * 8 + 2 * e_1] = __fmaf_rn(nkx, x2, rq * (g2 * wf[v_3 * 8 + 2 * e_1])); + out[v_3 * 8 + 2 * e_1 + 1] = __fmaf_rn(nkx, x3, rq * (g3 * wf[v_3 * 8 + 2 * e_1 + 1])); + } + int dcol = (v_3 * 96 + tir) * 8; + unsigned int _packed_0[4]; + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 0], (out + v_3 * 8)[_lp*2+1 + 0])); + _packed_0[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 2], (out + v_3 * 8)[_lp*2+1 + 2])); + _packed_0[_lp + 1] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 4], (out + v_3 * 8)[_lp*2+1 + 4])); + _packed_0[_lp + 2] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 6], (out + v_3 * 8)[_lp*2+1 + 6])); + _packed_0[_lp + 3] = *(uint32_t*)&_bf2; + } + reinterpret_cast(dx + (dx_base + (unsigned long long)dcol))[0] = reinterpret_cast(_packed_0)[0]; + } + } + r_cur = r_nxt; + g_step += 1; + if ((k + 1 & 3) == 0) { + #pragma unroll + for (int jkr = 0; jkr < 16; jkr++) { + float ykr = acc[jkr] - cmp[jkr]; + float tkr = hi[jkr] + ykr; + cmp[jkr] = tkr - hi[jkr] - ykr; + hi[jkr] = tkr; + acc[jkr] = 0.0f; + } + } + } + if ((n_steps_c & 3) != 0) { + #pragma unroll + for (int jkt = 0; jkt < 16; jkt++) { + float ykt = acc[jkt] - cmp[jkt]; + float tkt = hi[jkt] + ykt; + cmp[jkt] = tkt - hi[jkt] - ykt; + hi[jkt] = tkt; + acc[jkt] = 0.0f; + } + } + int entry_1_3 = n_cons % 5; + int _parity_dbar_3 = (int)((unsigned int)n_cons / 5 % 2); + mbarrier_wait(dbar_addr + (entry_1_3) * 8, _parity_dbar_3); + int word_1_3 = entry_1_3 * 3; + int w0_1 = desc[word_1_3]; + int w1_1 = desc[word_1_3 + 1]; + int w2_1 = desc[word_1_3 + 2]; + int _warp_redux_24; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_24) : "r"(w0_1)); + nd[0] = _warp_redux_24; + int _warp_redux_25; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_25) : "r"(w1_1)); + nd[1] = _warp_redux_25; + int _warp_redux_26; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_26) : "r"(w2_1)); + nd[2] = _warp_redux_26; + int cid_f = nd[0]; + unsigned long long pbase = (unsigned long long)cid_f * 1536; + __syncthreads(); + #pragma unroll + for (int jk = 0; jk < 16; jk++) { + acc[jk] = hi[jk] - cmp[jk]; + hi[jk] = 0.0f; + cmp[jk] = 0.0f; + } + #pragma unroll + for (int v_4 = 0; v_4 < 2; v_4++) { + int fcol = (v_4 * 96 + tir) * 8; + #pragma unroll + for (int e_2 = 0; e_2 < 8; e_2++) { + fold_sm[rg * 1536 + fcol + e_2] = acc[v_4 * 8 + e_2]; + } + } + __syncthreads(); + float fold_acc[4]; + #pragma unroll + for (int e_3 = 0; e_3 < 4; e_3++) { + fold_acc[e_3] = 0.0f; + } + #pragma unroll + for (int k2 = 0; k2 < 4; k2++) { + #pragma unroll + for (int e_4 = 0; e_4 < 4; e_4++) { + float d2 = fold_sm[k2 * 1536 + tid * 4 + e_4]; + fold_acc[e_4] = fold_acc[e_4] + d2; + } + } + int pcol = tid * 4; + { + float4 _v4 = make_float4(fold_acc[0 + 0], fold_acc[0 + 1], fold_acc[0 + 2], fold_acc[0 + 3]); + *reinterpret_cast(partial + pbase + (unsigned long long)pcol) = _v4; + } + #pragma unroll + for (int j2 = 0; j2 < 16; j2++) { + acc[j2] = 0.0f; + } + int entry_2 = (n_cons + 1) % 5; + int _parity_dbar_4 = (int)((unsigned int)(n_cons + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_2) * 8, _parity_dbar_4); + int word_2_2 = entry_2 * 3; + int w0_2_2 = desc[word_2_2]; + int w1_2 = desc[word_2_2 + 1]; + int w2_2 = desc[word_2_2 + 2]; + int _warp_redux_27; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_27) : "r"(w0_2_2)); + nd[0] = _warp_redux_27; + int _warp_redux_28; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_28) : "r"(w1_2)); + nd[1] = _warp_redux_28; + int _warp_redux_29; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_29) : "r"(w2_2)); + nd[2] = _warp_redux_29; + cb = nd[1]; + ce = nd[2]; + n_cons += 1; + } + + // Cleanup + __syncthreads(); +} + +} // extern "C" diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_binding.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_binding.cu new file mode 100644 index 00000000000..c8e20fc7a59 --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_binding.cu @@ -0,0 +1,146 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_det_train_16a5e0b1afd90f38ed44, for the device source that defines it. Do not edit. +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_rmsnorm_det_train_16a5e0b1afd90f38ed44(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0); + +namespace cake_binding_kernel_cake_rmsnorm_det_train_16a5e0b1afd90f38ed44 { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_det_train_16a5e0b1afd90f38ed44"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckRank(const TensorView& t, const char* name, int rank) { + TVM_FFI_CHECK(t.ndim() >= rank, ValueError) + << KERNEL_NAME << ": '" << name << "' needs at least " << rank << " dims, got " << t.ndim(); +} + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +void Run(TensorView a_g, TensorView a_x, TensorView a_w, TensorView a_r, TensorView a_dx, TensorView a_dw, TensorView a_partial, TensorView a_counters, TensorView a_g_h, int64_t a_T, int64_t a_n_chunks, int64_t a_rows_per_chunk, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_g.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_g, "g"); + check_dtype(a_g, DLDataType{kDLBfloat, 16, 1}, "g"); + check_cuda_tensor(a_x, "x"); + check_dtype(a_x, DLDataType{kDLBfloat, 16, 1}, "x"); + check_cuda_tensor(a_w, "w"); + check_dtype(a_w, DLDataType{kDLBfloat, 16, 1}, "w"); + check_cuda_tensor(a_r, "r"); + check_dtype(a_r, DLDataType{kDLFloat, 32, 1}, "r"); + check_cuda_tensor(a_dx, "dx"); + check_dtype(a_dx, DLDataType{kDLBfloat, 16, 1}, "dx"); + check_cuda_tensor(a_dw, "dw"); + check_dtype(a_dw, DLDataType{kDLFloat, 32, 1}, "dw"); + check_cuda_tensor(a_partial, "partial"); + check_dtype(a_partial, DLDataType{kDLFloat, 32, 1}, "partial"); + check_cuda_tensor(a_counters, "counters"); + check_dtype(a_counters, DLDataType{kDLUInt, 32, 1}, "counters"); + check_cuda_tensor(a_g_h, "g_h"); + check_dtype(a_g_h, DLDataType{kDLBfloat, 16, 1}, "g_h"); + TVM_FFI_CHECK(a_T >= -2147483648LL && a_T <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'T' = " << a_T + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_n_chunks >= -2147483648LL && a_n_chunks <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'n_chunks' = " << a_n_chunks + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_rows_per_chunk >= -2147483648LL && a_rows_per_chunk <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'rows_per_chunk' = " << a_rows_per_chunk + << " is outside the i32 range [-2147483648, 2147483647]"; + CheckRank(a_g, "g", 1); + CheckUnitInnerStride(a_g, "g"); + CheckRank(a_x, "x", 1); + CheckUnitInnerStride(a_x, "x"); + check_same_device(a_x, a_g, "x", "g"); + check_contiguous(a_w, "w"); + check_same_device(a_w, a_g, "w", "g"); + check_contiguous(a_r, "r"); + check_same_device(a_r, a_g, "r", "g"); + check_contiguous(a_dx, "dx"); + check_same_device(a_dx, a_g, "dx", "g"); + check_contiguous(a_dw, "dw"); + check_same_device(a_dw, a_g, "dw", "g"); + check_contiguous(a_partial, "partial"); + check_same_device(a_partial, a_g, "partial", "g"); + check_contiguous(a_counters, "counters"); + check_same_device(a_counters, a_g, "counters", "g"); + CheckRank(a_g_h, "g_h", 1); + CheckUnitInnerStride(a_g_h, "g_h"); + check_same_device(a_g_h, a_g, "g_h", "g"); + TVM_FFI_CHECK((a_n_chunks >= 1LL), ValueError) + << KERNEL_NAME << ": requires (n_chunks >= 1)"; + TVM_FFI_CHECK((a_rows_per_chunk >= 1LL), ValueError) + << KERNEL_NAME << ": requires (rows_per_chunk >= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_rmsnorm_det_train_16a5e0b1afd90f38ed44), cudaFuncAttributeMaxDynamicSharedMemorySize, 99840), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + void* v_g = DataPtr(a_g); + void* v_x = DataPtr(a_x); + void* v_w = DataPtr(a_w); + void* v_r = DataPtr(a_r); + void* v_dx = DataPtr(a_dx); + void* v_dw = DataPtr(a_dw); + void* v_partial = DataPtr(a_partial); + void* v_counters = DataPtr(a_counters); + void* v_g_h = DataPtr(a_g_h); + int32_t v_T = static_cast(a_T); + int32_t v_n_chunks = static_cast(a_n_chunks); + int32_t v_rows_per_chunk = static_cast(a_rows_per_chunk); + int64_t v_g_stride_p0 = a_g.stride(0); + int64_t v_x_stride_p0 = a_x.stride(0); + int64_t v_g_h_stride_p0 = a_g_h.stride(0); + void* kernel_args[] = {&v_g, &v_x, &v_w, &v_r, &v_dx, &v_dw, &v_partial, &v_counters, &v_g_h, &v_T, &v_n_chunks, &v_rows_per_chunk, &v_g_stride_p0, &v_x_stride_p0, &v_g_h_stride_p0}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(384, 1, 1); + config.dynamicSmemBytes = 99840; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_rmsnorm_det_train_16a5e0b1afd90f38ed44), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_rmsnorm_det_train_16a5e0b1afd90f38ed44 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_det_train_16a5e0b1afd90f38ed44::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_kernel.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_kernel.cu new file mode 100644 index 00000000000..9c4778db352 --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_16a5e0b1afd90f38ed44_kernel.cu @@ -0,0 +1,813 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define SMEM_STAGE_BUF_OFF 1024 +#define SMEM_STAGE_BUF_STAGE_BYTES 73728 +#define SMEM_STAGE_BUF_STRIDE 73728 +#define SMEM_ROW_FOLD_OFF 1024 +#define SMEM_ROW_FOLD_STAGE_BYTES 24576 +#define SMEM_ROW_FOLD_STRIDE 24576 +#define SMEM_RED_FOLD_OFF 25600 +#define SMEM_RED_FOLD_STAGE_BYTES 6144 +#define SMEM_RED_FOLD_STRIDE 6144 +#define SMEM_DOT_SLOTS_OFF 74752 +#define SMEM_DOT_SLOTS_STAGE_BYTES 96 +#define SMEM_DOT_SLOTS_STRIDE 96 +#define SMEM_PART_SM_OFF 74880 +#define SMEM_PART_SM_STAGE_BYTES 16 +#define SMEM_PART_SM_STRIDE 16 +#define SMEM_DESC_OFF 99584 +#define SMEM_DESC_STAGE_BYTES 60 +#define SMEM_DESC_STRIDE 60 +#define SMEM_RTAB_OFF 99712 +#define SMEM_RTAB_STAGE_BYTES 128 +#define SMEM_RTAB_STRIDE 128 +#define SMEM_FOLD_SM_OFF 75008 +#define SMEM_FOLD_SM_STAGE_BYTES 24576 +#define SMEM_FOLD_SM_STRIDE 24576 +#define SMEM_TOTAL 99840 +#define THREADS 384 +#define LAUNCH_MIN_BLOCKS 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_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"); +} + + + +extern "C" { + +__global__ __launch_bounds__(384, LAUNCH_MIN_BLOCKS) void +kernel_cake_rmsnorm_det_train_16a5e0b1afd90f38ed44(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define bars_addr (mbar_base + 0) + #define dbar_addr (mbar_base + 24) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + uint8_t* stage_buf = reinterpret_cast(smem_raw + 1024); + const int stage_buf_addr = smem + 1024; + float* row_fold = reinterpret_cast(smem_raw + 1024); + const int row_fold_addr = smem + 1024; + float* red_fold = reinterpret_cast(smem_raw + 25600); + const int red_fold_addr = smem + 25600; + float* dot_slots = reinterpret_cast(smem_raw + 74752); + const int dot_slots_addr = smem + 74752; + float* part_sm = reinterpret_cast(smem_raw + 74880); + const int part_sm_addr = smem + 74880; + int* desc = reinterpret_cast(smem_raw + 99584); + const int desc_addr = smem + 99584; + int* rtab = reinterpret_cast(smem_raw + 99712); + const int rtab_addr = smem + 99712; + float* fold_sm = reinterpret_cast(smem_raw + 75008); + const int fold_sm_addr = smem + 75008; + + // Mbarrier init (2 pipeline groups, 0 ordered-sequence groups, 8 barriers) + // Mbarriers at smem_raw[0..64) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // bars: 3 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + mbarrier_init(smem + 8, 1); + mbarrier_init(smem + 16, 1); + // dbar: 5 barriers, init_count=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); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncthreads(); + + // === Task calls (declaration order) === + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + long long g_stride = g_stride_p0; + long long x_stride = x_stride_p0; + long long gh_stride = g_h_stride_p0; + int c = bid; + int rg = tid / 96; + int tir = tid - rg * 96; + unsigned int claim[1]; + if (tid == 0) { + unsigned int _atomic_fetch_add_0 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_0; + int S = (T + 3) / 4; + int _min_1 = ((rows_per_chunk) < (S) ? (rows_per_chunk) : (S)); + int sl = ((_min_1) > (1) ? (_min_1) : (1)); + int share = (S + sl - 1) / sl; + int tail = (((share + 11) / 12) > (1) ? ((share + 11) / 12) : (1)); + int size = (share + 1) / 2; + int rem = S; + int base = 0; + int first = 0; + int nr = 0; + int walk = 1; + while (walk == 1) { + walk = 0; + if (nr < 5) { + if (size > tail) { + if (rem >= sl * size) { + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = size; + first += sl; + base += sl * size; + rem -= sl * size; + size = (size + 1) / 2; + nr += 1; + walk = 1; + } + } + } + } + int q = rem / sl; + int m = rem - q * sl; + if (m > 0) { + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = q + 1; + first += m; + base += m * (q + 1); + nr += 1; + } + if (q > 0) { + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = q; + first += sl - m; + base += (sl - m) * q; + nr += 1; + } + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = 0; + rtab[31] = nr + 1; + } + int ring = (int)stage_buf_addr; + int my_g = ring + rg * 3072 + tir * 16; + unsigned int ww[8]; + #pragma unroll + for (int v = 0; v < 2; v++) { + int wcol = (v * 96 + tir) * 8; + { + const uint4* _vptr_2 = reinterpret_cast(w + (unsigned long long)wcol); + uint4* _vdst_2 = reinterpret_cast(&ww[v * 4]); + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vdst_2[_blk] = _vptr_2[_blk]; + } + } + } + int rng[2]; + int nd[3]; + if (tid == 0) { + int cid0 = (int)claim[0]; + int nr_1 = rtab[31]; + int r_0 = 0; + int walk_1 = 1; + while (walk_1 == 1) { + walk_1 = 0; + if (nr_1 > r_0 + 2) { + if (cid0 >= rtab[(r_0 + 1) * 3]) { + r_0 += 1; + walk_1 = 1; + } + } + } + if (cid0 >= rtab[(nr_1 - 1) * 3]) { + rng[0] = T; + rng[1] = T; + } else { + int first_1 = rtab[r_0 * 3]; + int base_1 = rtab[r_0 * 3 + 1]; + int stages = rtab[r_0 * 3 + 2]; + rng[0] = (base_1 + (cid0 - first_1) * stages) * 4; + int _min_3 = ((rng[0] + stages * 4) < (T) ? (rng[0] + stages * 4) : (T)); + rng[1] = _min_3; + } + int entry = 0; + int word = entry * 3; + desc[word] = cid0; + desc[word + 1] = rng[0]; + desc[word + 2] = rng[1]; + mbarrier_arrive(dbar_addr + (entry) * 8); + if (rng[0] < T) { + unsigned int _atomic_fetch_add_4 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_4; + } else { + int nr_0 = rtab[31]; + int last = rtab[(nr_0 - 1) * 3] + num_bids - 1; + if (T == 0) { + last = 0; + } + if (cid0 == last) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } + } + int entry_1 = 0; + mbarrier_wait(dbar_addr + (entry_1) * 8, 0); + int word_1 = entry_1 * 3; + int w0 = desc[word_1]; + int w1 = desc[word_1 + 1]; + int w2 = desc[word_1 + 2]; + int _warp_redux_5; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_5) : "r"(w0)); + nd[0] = _warp_redux_5; + int _warp_redux_6; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_6) : "r"(w1)); + nd[1] = _warp_redux_6; + int _warp_redux_7; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_7) : "r"(w2)); + nd[2] = _warp_redux_7; + int c0 = nd[0]; + int cb = nd[1]; + int ce = nd[2]; + int ps[3]; + ps[0] = nd[1]; + ps[1] = nd[2]; + int desc_n = 0; + #pragma unroll + for (int s = 0; s < 3; s++) { + ps[2] = 0; + if (ps[0] < ps[1]) { + int first_2 = ps[0]; + int rows_k = ((ps[1] - first_2) < (4) ? (ps[1] - first_2) : (4)); + int sbase = ring + s * 24576; + if (tid == 0) { + mbarrier_arrive_expect_tx(bars_addr + (s) * 8, rows_k * 6144); + } + if (tir == 0) { + if (rg < rows_k) { + int row_g = first_2 + rg; + long long g_srcg = (long long)row_g * g_stride; + long long x_srcg = (long long)row_g * x_stride; + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase + rg * 3072), "l"(reinterpret_cast(g + g_srcg)), "r"((unsigned)(3072)), "r"(bars_addr + (s) * 8), + "l"(0x12F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase + (rg + 4) * 3072), "l"(reinterpret_cast(x + x_srcg)), "r"((unsigned)(3072)), "r"(bars_addr + (s) * 8), + "l"(0x12F0000000000000ULL) : "memory"); + } + } + ps[0] = ps[0] + 4; + ps[2] = (int)(ps[0] >= ps[1]); + if (ps[0] >= ps[1]) { + if (tid == 0) { + int cid = (int)claim[0]; + int nr_2 = rtab[31]; + int r_0_1 = 0; + int walk_2 = 1; + while (walk_2 == 1) { + walk_2 = 0; + if (nr_2 > r_0_1 + 2) { + if (cid >= rtab[(r_0_1 + 1) * 3]) { + r_0_1 += 1; + walk_2 = 1; + } + } + } + if (cid >= rtab[(nr_2 - 1) * 3]) { + rng[0] = T; + rng[1] = T; + } else { + int first_0 = rtab[r_0_1 * 3]; + int base_2 = rtab[r_0_1 * 3 + 1]; + int stages_1 = rtab[r_0_1 * 3 + 2]; + rng[0] = (base_2 + (cid - first_0) * stages_1) * 4; + int _min_8 = ((rng[0] + stages_1 * 4) < (T) ? (rng[0] + stages_1 * 4) : (T)); + rng[1] = _min_8; + } + int entry_1_1 = (desc_n + 1) % 5; + int word_2 = entry_1_1 * 3; + desc[word_2] = cid; + desc[word_2 + 1] = rng[0]; + desc[word_2 + 2] = rng[1]; + mbarrier_arrive(dbar_addr + (entry_1_1) * 8); + if (rng[0] < T) { + unsigned int _atomic_fetch_add_9 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_9; + } else { + int nr_0_1 = rtab[31]; + int last_1 = rtab[(nr_0_1 - 1) * 3] + num_bids - 1; + if (T == 0) { + last_1 = 0; + } + if (cid == last_1) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } + } + int entry_0 = (desc_n + 1) % 5; + int _parity_dbar = (int)((unsigned int)(desc_n + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_0) * 8, _parity_dbar); + int word_1_1 = entry_0 * 3; + int w0_2 = desc[word_1_1]; + int w1_3 = desc[word_1_1 + 1]; + int w2_4 = desc[word_1_1 + 2]; + int _warp_redux_10; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_10) : "r"(w0_2)); + nd[0] = _warp_redux_10; + int _warp_redux_11; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_11) : "r"(w1_3)); + nd[1] = _warp_redux_11; + int _warp_redux_12; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_12) : "r"(w2_4)); + nd[2] = _warp_redux_12; + ps[0] = nd[1]; + ps[1] = nd[2]; + } + } + int desc_step = ps[2]; + desc_n += desc_step; + } + float wf[16]; + #pragma unroll + for (int p = 0; p < 8; p++) { + unsigned int wwd = ww[p]; + wf[2 * p] = __uint_as_float(wwd << 16); + wf[2 * p + 1] = __uint_as_float(wwd & 4294901760u); + } + float acc[16]; + #pragma unroll + for (int j = 0; j < 16; j++) { + acc[j] = 0.0f; + } + float cmp[16]; + #pragma unroll + for (int jc = 0; jc < 16; jc++) { + cmp[jc] = 0.0f; + } + float hi[16]; + #pragma unroll + for (int jh = 0; jh < 16; jh++) { + hi[jh] = 0.0f; + } + int n_cons = 0; + int g_step = 0; + float r_cur = 0.0f; + if (cb < ce) { + int row0 = ((cb + rg) < (ce - 1) ? (cb + rg) : (ce - 1)); + r_cur = r[row0]; + } + unsigned int gw[8]; + unsigned int xw[8]; + float out[16]; + if (cb >= ce) { + if (c0 < n_chunks) { + unsigned long long pbase0 = (unsigned long long)c0 * 1536; + float fold0[4]; + #pragma unroll + for (int e0 = 0; e0 < 4; e0++) { + fold0[e0] = 0.0f; + } + int pcol0 = tid * 4; + { + float4 _v4 = make_float4(fold0[0 + 0], fold0[0 + 1], fold0[0 + 2], fold0[0 + 3]); + *reinterpret_cast(partial + pbase0 + (unsigned long long)pcol0) = _v4; + } + } + } + while (cb < ce) { + int n_steps_c = (ce - cb + 3) / 4; + #pragma unroll 1 + for (int k = 0; k < n_steps_c; k++) { + int slot = g_step % 3; + int par = (g_step & 1) * 12; + int row = cb + k * 4 + rg; + int row_n = ((row + 4) < (ce - 1) ? (row + 4) : (ce - 1)); + if (n_steps_c <= k + 1) { + int entry_0_1 = (n_cons + 1) % 5; + int _parity_dbar_1 = (int)((unsigned int)(n_cons + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_0_1) * 8, _parity_dbar_1); + int word_0 = entry_0_1 * 3; + int w0_0 = desc[word_0]; + int w1_0 = desc[word_0 + 1]; + int w2_0 = desc[word_0 + 2]; + int _warp_redux_13; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_13) : "r"(w0_0)); + nd[0] = _warp_redux_13; + int _warp_redux_14; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_14) : "r"(w1_0)); + nd[1] = _warp_redux_14; + int _warp_redux_15; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_15) : "r"(w2_0)); + nd[2] = _warp_redux_15; + row_n = cb; + if (nd[1] < T) { + int _min_16 = ((nd[1] + rg) < (nd[2] - 1) ? (nd[1] + rg) : (nd[2] - 1)); + row_n = _min_16; + } + } + float r_nxt = r[row_n]; + int _parity_bars = (int)((unsigned int)g_step / 3 % 2); + mbarrier_wait(bars_addr + (slot) * 8, _parity_bars); + int sg = my_g + slot * 24576; + #pragma unroll + for (int v_1 = 0; v_1 < 2; v_1++) { + __int128_t _smem_raw_17; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_17) : "r"(sg + v_1 * 1536)); + gw[v_1 * 4] = reinterpret_cast(&_smem_raw_17)[0]; + gw[(v_1 * 4) + 1] = reinterpret_cast(&_smem_raw_17)[1]; + gw[(v_1 * 4) + 2] = reinterpret_cast(&_smem_raw_17)[2]; + gw[(v_1 * 4) + 3] = reinterpret_cast(&_smem_raw_17)[3]; + __int128_t _smem_raw_18; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_18) : "r"(sg + (12288 + v_1 * 1536))); + xw[v_1 * 4] = reinterpret_cast(&_smem_raw_18)[0]; + xw[(v_1 * 4) + 1] = reinterpret_cast(&_smem_raw_18)[1]; + xw[(v_1 * 4) + 2] = reinterpret_cast(&_smem_raw_18)[2]; + xw[(v_1 * 4) + 3] = reinterpret_cast(&_smem_raw_18)[3]; + } + float dq = 0.0f; + #pragma unroll + for (int v_2 = 0; v_2 < 2; v_2++) { + #pragma unroll + for (int e = 0; e < 4; e++) { + unsigned int gwd = gw[v_2 * 4 + e]; + unsigned int xwd = xw[v_2 * 4 + e]; + float g0 = __uint_as_float(gwd << 16); + float g1 = __uint_as_float(gwd & 4294901760u); + float x0 = __uint_as_float(xwd << 16); + float x1 = __uint_as_float(xwd & 4294901760u); + dq = __fmaf_rn(g0 * wf[v_2 * 8 + 2 * e], x0, dq); + dq = __fmaf_rn(g1 * wf[v_2 * 8 + 2 * e + 1], x1, dq); + } + } + float dot = dq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + dot += __shfl_xor_sync(0xFFFFFFFF, dot, offset); + if (lane == 0) { + dot_slots[par + warp] = dot; + } + __syncthreads(); + ps[2] = 0; + if (ps[0] < ps[1]) { + int first_3 = ps[0]; + int rows_k_1 = ((ps[1] - first_3) < (4) ? (ps[1] - first_3) : (4)); + int sbase_1 = ring + g_step % 3 * 24576; + if (tid == 0) { + mbarrier_arrive_expect_tx(bars_addr + ((int)((unsigned int)g_step % 3)) * 8, rows_k_1 * 6144); + } + if (tir == 0) { + if (rg < rows_k_1) { + int row_g_1 = first_3 + rg; + long long g_srcg_1 = (long long)row_g_1 * g_stride; + long long x_srcg_1 = (long long)row_g_1 * x_stride; + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase_1 + rg * 3072), "l"(reinterpret_cast(g + g_srcg_1)), "r"((unsigned)(3072)), "r"(bars_addr + ((int)((unsigned int)g_step % 3)) * 8), + "l"(0x12F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase_1 + (rg + 4) * 3072), "l"(reinterpret_cast(x + x_srcg_1)), "r"((unsigned)(3072)), "r"(bars_addr + ((int)((unsigned int)g_step % 3)) * 8), + "l"(0x12F0000000000000ULL) : "memory"); + } + } + ps[0] = ps[0] + 4; + ps[2] = (int)(ps[0] >= ps[1]); + if (ps[0] >= ps[1]) { + if (tid == 0) { + int cid_1 = (int)claim[0]; + int nr_3 = rtab[31]; + int r_0_2 = 0; + int walk_3 = 1; + while (walk_3 == 1) { + walk_3 = 0; + if (nr_3 > r_0_2 + 2) { + if (cid_1 >= rtab[(r_0_2 + 1) * 3]) { + r_0_2 += 1; + walk_3 = 1; + } + } + } + if (cid_1 >= rtab[(nr_3 - 1) * 3]) { + rng[0] = T; + rng[1] = T; + } else { + int first_1_1 = rtab[r_0_2 * 3]; + int base_3 = rtab[r_0_2 * 3 + 1]; + int stages_2 = rtab[r_0_2 * 3 + 2]; + rng[0] = (base_3 + (cid_1 - first_1_1) * stages_2) * 4; + int _min_19 = ((rng[0] + stages_2 * 4) < (T) ? (rng[0] + stages_2 * 4) : (T)); + rng[1] = _min_19; + } + int entry_1_2 = (desc_n + 1) % 5; + int word_2_1 = entry_1_2 * 3; + desc[word_2_1] = cid_1; + desc[word_2_1 + 1] = rng[0]; + desc[word_2_1 + 2] = rng[1]; + mbarrier_arrive(dbar_addr + (entry_1_2) * 8); + if (rng[0] < T) { + unsigned int _atomic_fetch_add_20 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_20; + } else { + int nr_0_2 = rtab[31]; + int last_2 = rtab[(nr_0_2 - 1) * 3] + num_bids - 1; + if (T == 0) { + last_2 = 0; + } + if (cid_1 == last_2) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } + } + int entry_0_2 = (desc_n + 1) % 5; + int _parity_dbar_2 = (int)((unsigned int)(desc_n + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_0_2) * 8, _parity_dbar_2); + int word_1_2 = entry_0_2 * 3; + int w0_2_1 = desc[word_1_2]; + int w1_3_1 = desc[word_1_2 + 1]; + int w2_4_1 = desc[word_1_2 + 2]; + int _warp_redux_21; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_21) : "r"(w0_2_1)); + nd[0] = _warp_redux_21; + int _warp_redux_22; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_22) : "r"(w1_3_1)); + nd[1] = _warp_redux_22; + int _warp_redux_23; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_23) : "r"(w2_4_1)); + nd[2] = _warp_redux_23; + ps[0] = nd[1]; + ps[1] = nd[2]; + } + } + int desc_step_1 = ps[2]; + desc_n += desc_step_1; + float tot = 0.0f; + #pragma unroll + for (int kk = 0; kk < 3; kk++) { + float d = dot_slots[par + rg * 3 + kk]; + tot += d; + } + float rq = r_cur; + float cq = __fdiv_rn(tot, 1536.0f); + float kx = rq * rq * rq * cq; + float nkx = -kx; + if (row < ce) { + unsigned long long dx_base = (unsigned long long)row * 1536; + #pragma unroll + for (int v_3 = 0; v_3 < 2; v_3++) { + #pragma unroll + for (int e_1 = 0; e_1 < 4; e_1++) { + unsigned int gwd2 = gw[v_3 * 4 + e_1]; + unsigned int xwd2 = xw[v_3 * 4 + e_1]; + float g2 = __uint_as_float(gwd2 << 16); + float g3 = __uint_as_float(gwd2 & 4294901760u); + float x2 = __uint_as_float(xwd2 << 16); + float x3 = __uint_as_float(xwd2 & 4294901760u); + acc[v_3 * 8 + 2 * e_1] = __fmaf_rn(g2 * x2, rq, acc[v_3 * 8 + 2 * e_1]); + acc[v_3 * 8 + 2 * e_1 + 1] = __fmaf_rn(g3 * x3, rq, acc[v_3 * 8 + 2 * e_1 + 1]); + out[v_3 * 8 + 2 * e_1] = __fmaf_rn(nkx, x2, rq * (g2 * wf[v_3 * 8 + 2 * e_1])); + out[v_3 * 8 + 2 * e_1 + 1] = __fmaf_rn(nkx, x3, rq * (g3 * wf[v_3 * 8 + 2 * e_1 + 1])); + } + int dcol = (v_3 * 96 + tir) * 8; + unsigned int _packed_0[4]; + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 0], (out + v_3 * 8)[_lp*2+1 + 0])); + _packed_0[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 2], (out + v_3 * 8)[_lp*2+1 + 2])); + _packed_0[_lp + 1] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 4], (out + v_3 * 8)[_lp*2+1 + 4])); + _packed_0[_lp + 2] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 6], (out + v_3 * 8)[_lp*2+1 + 6])); + _packed_0[_lp + 3] = *(uint32_t*)&_bf2; + } + reinterpret_cast(dx + (dx_base + (unsigned long long)dcol))[0] = reinterpret_cast(_packed_0)[0]; + } + } + r_cur = r_nxt; + g_step += 1; + if ((k + 1 & 15) == 0) { + #pragma unroll + for (int jkr = 0; jkr < 16; jkr++) { + float ykr = acc[jkr] - cmp[jkr]; + float tkr = hi[jkr] + ykr; + cmp[jkr] = tkr - hi[jkr] - ykr; + hi[jkr] = tkr; + acc[jkr] = 0.0f; + } + } + } + if ((n_steps_c & 15) != 0) { + #pragma unroll + for (int jkt = 0; jkt < 16; jkt++) { + float ykt = acc[jkt] - cmp[jkt]; + float tkt = hi[jkt] + ykt; + cmp[jkt] = tkt - hi[jkt] - ykt; + hi[jkt] = tkt; + acc[jkt] = 0.0f; + } + } + int entry_1_3 = n_cons % 5; + int _parity_dbar_3 = (int)((unsigned int)n_cons / 5 % 2); + mbarrier_wait(dbar_addr + (entry_1_3) * 8, _parity_dbar_3); + int word_1_3 = entry_1_3 * 3; + int w0_1 = desc[word_1_3]; + int w1_1 = desc[word_1_3 + 1]; + int w2_1 = desc[word_1_3 + 2]; + int _warp_redux_24; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_24) : "r"(w0_1)); + nd[0] = _warp_redux_24; + int _warp_redux_25; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_25) : "r"(w1_1)); + nd[1] = _warp_redux_25; + int _warp_redux_26; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_26) : "r"(w2_1)); + nd[2] = _warp_redux_26; + int cid_f = nd[0]; + unsigned long long pbase = (unsigned long long)cid_f * 1536; + #pragma unroll + for (int jk = 0; jk < 16; jk++) { + acc[jk] = hi[jk] - cmp[jk]; + hi[jk] = 0.0f; + cmp[jk] = 0.0f; + } + #pragma unroll + for (int v_4 = 0; v_4 < 2; v_4++) { + int fcol = (v_4 * 96 + tir) * 8; + #pragma unroll + for (int e_2 = 0; e_2 < 8; e_2++) { + fold_sm[rg * 1536 + fcol + e_2] = acc[v_4 * 8 + e_2]; + } + } + __syncthreads(); + float fold_acc[4]; + #pragma unroll + for (int e_3 = 0; e_3 < 4; e_3++) { + fold_acc[e_3] = 0.0f; + } + #pragma unroll + for (int k2 = 0; k2 < 4; k2++) { + #pragma unroll + for (int e_4 = 0; e_4 < 4; e_4++) { + float d2 = fold_sm[k2 * 1536 + tid * 4 + e_4]; + fold_acc[e_4] = fold_acc[e_4] + d2; + } + } + int pcol = tid * 4; + { + float4 _v4 = make_float4(fold_acc[0 + 0], fold_acc[0 + 1], fold_acc[0 + 2], fold_acc[0 + 3]); + *reinterpret_cast(partial + pbase + (unsigned long long)pcol) = _v4; + } + #pragma unroll + for (int j2 = 0; j2 < 16; j2++) { + acc[j2] = 0.0f; + } + int entry_2 = (n_cons + 1) % 5; + int _parity_dbar_4 = (int)((unsigned int)(n_cons + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_2) * 8, _parity_dbar_4); + int word_2_2 = entry_2 * 3; + int w0_2_2 = desc[word_2_2]; + int w1_2 = desc[word_2_2 + 1]; + int w2_2 = desc[word_2_2 + 2]; + int _warp_redux_27; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_27) : "r"(w0_2_2)); + nd[0] = _warp_redux_27; + int _warp_redux_28; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_28) : "r"(w1_2)); + nd[1] = _warp_redux_28; + int _warp_redux_29; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_29) : "r"(w2_2)); + nd[2] = _warp_redux_29; + cb = nd[1]; + ce = nd[2]; + n_cons += 1; + } + + // Cleanup + __syncthreads(); +} + +} // extern "C" diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_affa1b95f68bab8c411b_binding.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_affa1b95f68bab8c411b_binding.cu new file mode 100644 index 00000000000..40fccb20a3b --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_affa1b95f68bab8c411b_binding.cu @@ -0,0 +1,80 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_det_train_affa1b95f68bab8c411b, for the device source that defines it. Do not edit. +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_rmsnorm_det_train_affa1b95f68bab8c411b(float* __restrict__ partial, float* __restrict__ dw, int n_chunks); + +namespace cake_binding_kernel_cake_rmsnorm_det_train_affa1b95f68bab8c411b { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_det_train_affa1b95f68bab8c411b"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void Run(TensorView a_partial, TensorView a_dw, int64_t a_n_chunks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_partial.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_partial, "partial"); + check_dtype(a_partial, DLDataType{kDLFloat, 32, 1}, "partial"); + check_cuda_tensor(a_dw, "dw"); + check_dtype(a_dw, DLDataType{kDLFloat, 32, 1}, "dw"); + TVM_FFI_CHECK(a_n_chunks >= -2147483648LL && a_n_chunks <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'n_chunks' = " << a_n_chunks + << " is outside the i32 range [-2147483648, 2147483647]"; + check_contiguous(a_partial, "partial"); + check_contiguous(a_dw, "dw"); + check_same_device(a_dw, a_partial, "dw", "partial"); + TVM_FFI_CHECK((a_n_chunks >= 1LL), ValueError) + << KERNEL_NAME << ": requires (n_chunks >= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + void* v_partial = DataPtr(a_partial); + void* v_dw = DataPtr(a_dw); + int32_t v_n_chunks = static_cast(a_n_chunks); + void* kernel_args[] = {&v_partial, &v_dw, &v_n_chunks}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(128, 1, 1); + config.dynamicSmemBytes = 2048; + config.stream = stream; + cudaLaunchAttribute attrs[1] = {}; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = 1; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_rmsnorm_det_train_affa1b95f68bab8c411b), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_rmsnorm_det_train_affa1b95f68bab8c411b + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_det_train_affa1b95f68bab8c411b::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_affa1b95f68bab8c411b_kernel.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_affa1b95f68bab8c411b_kernel.cu new file mode 100644 index 00000000000..d4e06279f86 --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_affa1b95f68bab8c411b_kernel.cu @@ -0,0 +1,178 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define SMEM_RED_FOLD_OFF 0 +#define SMEM_RED_FOLD_STAGE_BYTES 2048 +#define SMEM_RED_FOLD_STRIDE 2048 +#define SMEM_TOTAL 2048 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_cake_rmsnorm_det_train_affa1b95f68bab8c411b(float* __restrict__ partial, float* __restrict__ dw, int n_chunks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + float* red_fold = reinterpret_cast(smem_raw + 0); + const int red_fold_addr = smem + 0; + + // === Task calls (declaration order) === + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + int s = bid; + int rsub = tid / 4; + int cg = tid - rsub * 4; + int col = s * 4 + cg; + float a[1]; + float c[1]; + #pragma unroll + for (int e = 0; e < 1; e++) { + a[e] = 0.0f; + c[e] = 0.0f; + } + int n_trips = (n_chunks + 31) / 32; + int n_waves = (n_trips + 23) / 24; + float cur[24]; + #pragma unroll + for (int i = 0; i < 24; i++) { + cur[i] = 0.0f; + } + #pragma unroll + for (int d = 0; d < 24; d++) { + int cc = rsub + d * 32; + if (cc < n_chunks) { + { + cur[d] = *reinterpret_cast(partial + ((unsigned long long)cc * 256 + (unsigned long long)col)); + } + } + } + for (int w = 0; w < n_waves; w++) { + float nxt[24]; + #pragma unroll + for (int i_1 = 0; i_1 < 24; i_1++) { + nxt[i_1] = 0.0f; + } + #pragma unroll + for (int d_1 = 0; d_1 < 24; d_1++) { + int cc_1 = rsub + ((w + 1) * 24 + d_1) * 32; + if (cc_1 < n_chunks) { + { + nxt[d_1] = *reinterpret_cast(partial + ((unsigned long long)cc_1 * 256 + (unsigned long long)col)); + } + } + } + #pragma unroll + for (int d_2 = 0; d_2 < 24; d_2++) { + int cc_2 = rsub + (w * 24 + d_2) * 32; + if (cc_2 < n_chunks) { + #pragma unroll + for (int e_1 = 0; e_1 < 1; e_1++) { + float y = __fmaf_rn(cur[d_2 + e_1], 1.0f, -c[e_1]); + float t = a[e_1] + y; + c[e_1] = t - a[e_1] - y; + a[e_1] = t; + } + } + } + #pragma unroll + for (int i_2 = 0; i_2 < 24; i_2++) { + cur[i_2] = nxt[i_2]; + } + } + #pragma unroll + for (int e_2 = 0; e_2 < 1; e_2++) { + red_fold[rsub * 4 + cg + e_2] = a[e_2] - c[e_2]; + } + __syncthreads(); + if (tid < 4) { + float o4[1]; + float oc4[1]; + #pragma unroll + for (int e_3 = 0; e_3 < 1; e_3++) { + o4[e_3] = 0.0f; + oc4[e_3] = 0.0f; + } + #pragma unroll + for (int k = 0; k < 32; k++) { + #pragma unroll + for (int e_4 = 0; e_4 < 1; e_4++) { + float d_3 = red_fold[k * 4 + tid + e_4]; + float y_1 = __fmaf_rn(d_3, 1.0f, -oc4[e_4]); + float t_1 = o4[e_4] + y_1; + oc4[e_4] = t_1 - o4[e_4] - y_1; + o4[e_4] = t_1; + } + } + #pragma unroll + for (int e_5 = 0; e_5 < 1; e_5++) { + o4[e_5] = o4[e_5] - oc4[e_5]; + } + int ocol = s * 4 + tid; + *(reinterpret_cast(dw) + ((unsigned long long)ocol)) = o4[0]; + } + __syncthreads(); +} + +} // extern "C" diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_binding.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_binding.cu new file mode 100644 index 00000000000..680561bf4b8 --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_binding.cu @@ -0,0 +1,80 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_det_train_b4b63d5b5028c7a60333, for the device source that defines it. Do not edit. +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_rmsnorm_det_train_b4b63d5b5028c7a60333(float* __restrict__ partial, float* __restrict__ dw, int n_chunks); + +namespace cake_binding_kernel_cake_rmsnorm_det_train_b4b63d5b5028c7a60333 { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_det_train_b4b63d5b5028c7a60333"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void Run(TensorView a_partial, TensorView a_dw, int64_t a_n_chunks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_partial.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_partial, "partial"); + check_dtype(a_partial, DLDataType{kDLFloat, 32, 1}, "partial"); + check_cuda_tensor(a_dw, "dw"); + check_dtype(a_dw, DLDataType{kDLFloat, 32, 1}, "dw"); + TVM_FFI_CHECK(a_n_chunks >= -2147483648LL && a_n_chunks <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'n_chunks' = " << a_n_chunks + << " is outside the i32 range [-2147483648, 2147483647]"; + check_contiguous(a_partial, "partial"); + check_contiguous(a_dw, "dw"); + check_same_device(a_dw, a_partial, "dw", "partial"); + TVM_FFI_CHECK((a_n_chunks >= 1LL), ValueError) + << KERNEL_NAME << ": requires (n_chunks >= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + void* v_partial = DataPtr(a_partial); + void* v_dw = DataPtr(a_dw); + int32_t v_n_chunks = static_cast(a_n_chunks); + void* kernel_args[] = {&v_partial, &v_dw, &v_n_chunks}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(128, 1, 1); + config.dynamicSmemBytes = 2048; + config.stream = stream; + cudaLaunchAttribute attrs[1] = {}; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = 1; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_rmsnorm_det_train_b4b63d5b5028c7a60333), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_rmsnorm_det_train_b4b63d5b5028c7a60333 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_det_train_b4b63d5b5028c7a60333::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_kernel.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_kernel.cu new file mode 100644 index 00000000000..723a5b203fa --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_b4b63d5b5028c7a60333_kernel.cu @@ -0,0 +1,199 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define SMEM_RED_FOLD_OFF 0 +#define SMEM_RED_FOLD_STAGE_BYTES 2048 +#define SMEM_RED_FOLD_STRIDE 2048 +#define SMEM_TOTAL 2048 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_cake_rmsnorm_det_train_b4b63d5b5028c7a60333(float* __restrict__ partial, float* __restrict__ dw, int n_chunks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + float* red_fold = reinterpret_cast(smem_raw + 0); + const int red_fold_addr = smem + 0; + + // === Task calls (declaration order) === + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + int s = bid; + int rsub = tid / 4; + int cg = tid - rsub * 4; + int col = s * 4 + cg; + float a[1]; + float c[1]; + #pragma unroll + for (int e = 0; e < 1; e++) { + a[e] = 0.0f; + c[e] = 0.0f; + } + int n_trips = (n_chunks + 31) / 32; + int n_waves = (n_trips + 23) / 24; + float cur[24]; + #pragma unroll + for (int i = 0; i < 24; i++) { + cur[i] = 0.0f; + } + #pragma unroll + for (int d = 0; d < 24; d++) { + int cc = rsub + d * 32; + if (cc < n_chunks) { + { + cur[d] = *reinterpret_cast(partial + ((unsigned long long)cc * 256 + (unsigned long long)col)); + } + } + } + for (int w = 0; w < n_waves; w++) { + float nxt[24]; + #pragma unroll + for (int i_1 = 0; i_1 < 24; i_1++) { + nxt[i_1] = 0.0f; + } + #pragma unroll + for (int d_1 = 0; d_1 < 24; d_1++) { + int cc_1 = rsub + ((w + 1) * 24 + d_1) * 32; + if (cc_1 < n_chunks) { + { + nxt[d_1] = *reinterpret_cast(partial + ((unsigned long long)cc_1 * 256 + (unsigned long long)col)); + } + } + } + #pragma unroll + for (int d_2 = 0; d_2 < 24; d_2++) { + int cc_2 = rsub + (w * 24 + d_2) * 32; + if (cc_2 < n_chunks) { + #pragma unroll + for (int e_1 = 0; e_1 < 1; e_1++) { + float s_0 = a[e_1] + cur[d_2 + e_1]; + float bb = s_0 - a[e_1]; + c[e_1] = c[e_1] + (a[e_1] - (s_0 - bb) + (cur[d_2 + e_1] - bb)); + a[e_1] = s_0; + } + } + } + #pragma unroll + for (int i_2 = 0; i_2 < 24; i_2++) { + cur[i_2] = nxt[i_2]; + } + } + float o4[1]; + float oc4[1]; + #pragma unroll + for (int e_2 = 0; e_2 < 1; e_2++) { + o4[e_2] = 0.0f; + oc4[e_2] = 0.0f; + } + #pragma unroll + for (int e_3 = 0; e_3 < 1; e_3++) { + red_fold[rsub * 4 + cg + e_3] = a[e_3]; + } + __syncthreads(); + if (tid < 4) { + #pragma unroll + for (int k = 0; k < 32; k++) { + #pragma unroll + for (int e_4 = 0; e_4 < 1; e_4++) { + float d_3 = red_fold[k * 4 + tid + e_4]; + float s_1 = o4[e_4] + d_3; + float bb_1 = s_1 - o4[e_4]; + oc4[e_4] = oc4[e_4] + (o4[e_4] - (s_1 - bb_1) + (d_3 - bb_1)); + o4[e_4] = s_1; + } + } + } + __syncthreads(); + #pragma unroll + for (int e_5 = 0; e_5 < 1; e_5++) { + red_fold[rsub * 4 + cg + e_5] = c[e_5]; + } + __syncthreads(); + if (tid < 4) { + #pragma unroll + for (int k_1 = 0; k_1 < 32; k_1++) { + #pragma unroll + for (int e_6 = 0; e_6 < 1; e_6++) { + float dn = red_fold[k_1 * 4 + tid + e_6]; + oc4[e_6] = oc4[e_6] + dn; + } + } + #pragma unroll + for (int e_7 = 0; e_7 < 1; e_7++) { + o4[e_7] = o4[e_7] + oc4[e_7]; + } + int ocol = s * 4 + tid; + *(reinterpret_cast(dw) + ((unsigned long long)ocol)) = o4[0]; + } + __syncthreads(); +} + +} // extern "C" diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_binding.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_binding.cu new file mode 100644 index 00000000000..6bc9b71560a --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_binding.cu @@ -0,0 +1,137 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b, for the device source that defines it. Do not edit. +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, float* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0); + +namespace cake_binding_kernel_cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckRank(const TensorView& t, const char* name, int rank) { + TVM_FFI_CHECK(t.ndim() >= rank, ValueError) + << KERNEL_NAME << ": '" << name << "' needs at least " << rank << " dims, got " << t.ndim(); +} + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +void Run(TensorView a_g, TensorView a_x, TensorView a_w, TensorView a_r, TensorView a_dx, TensorView a_dw, TensorView a_partial, TensorView a_counters, TensorView a_g_h, int64_t a_T, int64_t a_n_chunks, int64_t a_rows_per_chunk, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_g.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_g, "g"); + check_dtype(a_g, DLDataType{kDLBfloat, 16, 1}, "g"); + check_cuda_tensor(a_x, "x"); + check_dtype(a_x, DLDataType{kDLBfloat, 16, 1}, "x"); + check_cuda_tensor(a_w, "w"); + check_dtype(a_w, DLDataType{kDLFloat, 32, 1}, "w"); + check_cuda_tensor(a_r, "r"); + check_dtype(a_r, DLDataType{kDLFloat, 32, 1}, "r"); + check_cuda_tensor(a_dx, "dx"); + check_dtype(a_dx, DLDataType{kDLBfloat, 16, 1}, "dx"); + check_cuda_tensor(a_dw, "dw"); + check_dtype(a_dw, DLDataType{kDLFloat, 32, 1}, "dw"); + check_cuda_tensor(a_partial, "partial"); + check_dtype(a_partial, DLDataType{kDLFloat, 32, 1}, "partial"); + check_cuda_tensor(a_counters, "counters"); + check_dtype(a_counters, DLDataType{kDLUInt, 32, 1}, "counters"); + check_cuda_tensor(a_g_h, "g_h"); + check_dtype(a_g_h, DLDataType{kDLBfloat, 16, 1}, "g_h"); + TVM_FFI_CHECK(a_T >= -2147483648LL && a_T <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'T' = " << a_T + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_n_chunks >= -2147483648LL && a_n_chunks <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'n_chunks' = " << a_n_chunks + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_rows_per_chunk >= -2147483648LL && a_rows_per_chunk <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'rows_per_chunk' = " << a_rows_per_chunk + << " is outside the i32 range [-2147483648, 2147483647]"; + CheckRank(a_g, "g", 1); + CheckUnitInnerStride(a_g, "g"); + CheckRank(a_x, "x", 1); + CheckUnitInnerStride(a_x, "x"); + check_same_device(a_x, a_g, "x", "g"); + check_contiguous(a_w, "w"); + check_same_device(a_w, a_g, "w", "g"); + check_contiguous(a_r, "r"); + check_same_device(a_r, a_g, "r", "g"); + check_contiguous(a_dx, "dx"); + check_same_device(a_dx, a_g, "dx", "g"); + check_contiguous(a_dw, "dw"); + check_same_device(a_dw, a_g, "dw", "g"); + check_contiguous(a_partial, "partial"); + check_same_device(a_partial, a_g, "partial", "g"); + check_contiguous(a_counters, "counters"); + check_same_device(a_counters, a_g, "counters", "g"); + CheckRank(a_g_h, "g_h", 1); + CheckUnitInnerStride(a_g_h, "g_h"); + check_same_device(a_g_h, a_g, "g_h", "g"); + TVM_FFI_CHECK((a_n_chunks >= 1LL), ValueError) + << KERNEL_NAME << ": requires (n_chunks >= 1)"; + TVM_FFI_CHECK((a_rows_per_chunk >= 1LL), ValueError) + << KERNEL_NAME << ": requires (rows_per_chunk >= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + void* v_g = DataPtr(a_g); + void* v_x = DataPtr(a_x); + void* v_w = DataPtr(a_w); + void* v_r = DataPtr(a_r); + void* v_dx = DataPtr(a_dx); + void* v_dw = DataPtr(a_dw); + void* v_partial = DataPtr(a_partial); + void* v_counters = DataPtr(a_counters); + void* v_g_h = DataPtr(a_g_h); + int32_t v_T = static_cast(a_T); + int32_t v_n_chunks = static_cast(a_n_chunks); + int32_t v_rows_per_chunk = static_cast(a_rows_per_chunk); + int64_t v_g_stride_p0 = a_g.stride(0); + int64_t v_x_stride_p0 = a_x.stride(0); + int64_t v_g_h_stride_p0 = a_g_h.stride(0); + void* kernel_args[] = {&v_g, &v_x, &v_w, &v_r, &v_dx, &v_dw, &v_partial, &v_counters, &v_g_h, &v_T, &v_n_chunks, &v_rows_per_chunk, &v_g_stride_p0, &v_x_stride_p0, &v_g_h_stride_p0}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(128, 1, 1); + config.dynamicSmemBytes = 43648; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_kernel.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_kernel.cu new file mode 100644 index 00000000000..a034aa0e60f --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b_kernel.cu @@ -0,0 +1,780 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define SMEM_DOT_SLOTS_OFF 0 +#define SMEM_DOT_SLOTS_STAGE_BYTES 32 +#define SMEM_DOT_SLOTS_STRIDE 32 +#define SMEM_FLAGS_OFF 128 +#define SMEM_FLAGS_STAGE_BYTES 32 +#define SMEM_FLAGS_STRIDE 32 +#define SMEM_RED_FOLD_OFF 256 +#define SMEM_RED_FOLD_STAGE_BYTES 2048 +#define SMEM_RED_FOLD_STRIDE 2048 +#define SMEM_ROW_FOLD_OFF 2304 +#define SMEM_ROW_FOLD_STAGE_BYTES 4096 +#define SMEM_ROW_FOLD_STRIDE 4096 +#define SMEM_COMP_SMEM_OFF 6400 +#define SMEM_COMP_SMEM_STAGE_BYTES 128 +#define SMEM_COMP_SMEM_STRIDE 128 +#define SMEM_W_SM_OFF 6528 +#define SMEM_W_SM_STAGE_BYTES 16 +#define SMEM_W_SM_STRIDE 16 +#define SMEM_PART_SM_OFF 6656 +#define SMEM_PART_SM_STAGE_BYTES 16 +#define SMEM_PART_SM_STRIDE 16 +#define SMEM_RING_OFF 6784 +#define SMEM_RING_STAGE_BYTES 36864 +#define SMEM_RING_STRIDE 36864 +#define SMEM_TOTAL 43648 +#define THREADS 128 +#define LAUNCH_MIN_BLOCKS 5 + +#include + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.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 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; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__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; +} + +extern "C" { + +__global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void +kernel_cake_rmsnorm_det_train_c4c89f0d04b7b3a38d0b(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, float* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + float* dot_slots = reinterpret_cast(smem_raw + 0); + const int dot_slots_addr = smem + 0; + int* flags = reinterpret_cast(smem_raw + 128); + const int flags_addr = smem + 128; + float* red_fold = reinterpret_cast(smem_raw + 256); + const int red_fold_addr = smem + 256; + float* row_fold = reinterpret_cast(smem_raw + 2304); + const int row_fold_addr = smem + 2304; + float* comp_smem = reinterpret_cast(smem_raw + 6400); + const int comp_smem_addr = smem + 6400; + unsigned int* w_sm = reinterpret_cast(smem_raw + 6528); + const int w_sm_addr = smem + 6528; + float* part_sm = reinterpret_cast(smem_raw + 6656); + const int part_sm_addr = smem + 6656; + uint8_t* ring = reinterpret_cast(smem_raw + 6784); + const int ring_addr = smem + 6784; + + // === Task calls (declaration order) === + asm volatile("griddepcontrol.wait;" ::: "memory"); + unsigned int nxt = 0; + if (tid == 0) { + unsigned int _atomic_fetch_add_0 = atomicAdd(&counters[0], 1); + nxt = _atomic_fetch_add_0; + } + long long g_stride = g_stride_p0; + long long x_stride = x_stride_p0; + long long gh_stride = g_h_stride_p0; + int c = 0; + int rg = warp; + int tir = tid - rg * 32; + int row_begin = 0; + int row_end = 0; + int n_rows = 0; + int n_iters = 0; + int last_row = 0; + float wc[8]; + #pragma unroll + for (int v = 0; v < 1; v++) { + int wcol = (v * 32 + tir) * 8; + { + unsigned _ldv8_1_0; + unsigned _ldv8_1_1; + unsigned _ldv8_1_2; + unsigned _ldv8_1_3; + unsigned _ldv8_1_4; + unsigned _ldv8_1_5; + unsigned _ldv8_1_6; + unsigned _ldv8_1_7; + asm volatile( + "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_1_0), "=r"(_ldv8_1_1), "=r"(_ldv8_1_2), "=r"(_ldv8_1_3), "=r"(_ldv8_1_4), "=r"(_ldv8_1_5), "=r"(_ldv8_1_6), "=r"(_ldv8_1_7) : "l"((const void*)(w + ((unsigned long long)wcol))) : "memory"); + wc[v * 8 + 0] = __uint_as_float(_ldv8_1_0); + wc[v * 8 + 1] = __uint_as_float(_ldv8_1_1); + wc[v * 8 + 2] = __uint_as_float(_ldv8_1_2); + wc[v * 8 + 3] = __uint_as_float(_ldv8_1_3); + wc[v * 8 + 4] = __uint_as_float(_ldv8_1_4); + wc[v * 8 + 5] = __uint_as_float(_ldv8_1_5); + wc[v * 8 + 6] = __uint_as_float(_ldv8_1_6); + wc[v * 8 + 7] = __uint_as_float(_ldv8_1_7); + } + } + float acc[8]; + float comp[8]; + unsigned int compw[4]; + if (tid == 0) { + flags[1] = (int)nxt; + } + __syncthreads(); + c = flags[1]; + #pragma unroll 1 + for (int claim_it = 0; claim_it < n_chunks; claim_it++) { + if (c >= n_chunks) { + break; + } + if (tid == 0) { + unsigned int _atomic_fetch_add_2 = atomicAdd(&counters[0], 1); + nxt = _atomic_fetch_add_2; + } + int stages = (T + 3) / 4; + int sl = rows_per_chunk; + if (sl > stages) { + sl = stages; + } + if (sl < 1) { + sl = 1; + } + int share = (stages + sl - 1) / sl; + int tail = (share + 3) / 4; + if (tail < 1) { + tail = 1; + } + int rem = stages; + int size = (share + 1) / 2; + int first = 0; + int base = 0; + int rows = 0; + int halving = 1; + int begin_it = 0; + #pragma unroll 1 + for (int k = 0; k < 5; k++) { + if (halving != 0) { + if (size > tail) { + if (rem >= sl * size) { + if (c < first + sl) { + rows = size; + begin_it = base + (c - first); + halving = 0; + } else { + first = first + sl; + base = base + sl * size; + rem = rem - sl * size; + size = (size + 1) / 2; + } + } else { + halving = 0; + } + } else { + halving = 0; + } + } + } + if (rows == 0) { + int quot = rem / sl; + int extra = rem - quot * sl; + if (c < first + extra) { + rows = quot + 1; + begin_it = base + (c - first); + } else if (c < first + sl) { + rows = quot; + begin_it = base + (c - first); + } else { + begin_it = stages; + } + } + row_begin = begin_it * 4; + row_end = row_begin + rows * 4; + if (row_begin > T) { + row_begin = T; + } + if (row_end > T) { + row_end = T; + } + n_rows = row_end - row_begin; + n_iters = (n_rows + 3) / 4; + last_row = row_end - 1; + int rstep = sl * 4; + row_end = T; + n_iters = rows; + last_row = T - 1; + #pragma unroll + for (int j = 0; j < 8; j++) { + acc[j] = 0.0f; + comp[j] = 0.0f; + } + unsigned int gwA[4]; + unsigned int xwA[4]; + unsigned int ghwA[4]; + float rrA[1]; + unsigned int gwB[4]; + unsigned int xwB[4]; + unsigned int ghwB[4]; + float rrB[1]; + if (n_iters > 0) { + int row0 = ((row_begin + rg) < (last_row) ? (row_begin + rg) : (last_row)); + unsigned int sg = (unsigned int)ring_addr + (unsigned int)(rg * 9216) + (unsigned int)(tir * 16); + unsigned int sr = (unsigned int)ring_addr + (unsigned int)(rg * 9216 + 1024) + (unsigned int)(tir * 4); + #pragma unroll + for (int p = 0; p < 3; p++) { + int rowq_ld = ((row_begin + p * rstep + rg) < (last_row) ? (row_begin + p * rstep + rg) : (last_row)); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.cg.shared::cta.global [%1], [%2], 16;\n\t" + "}" + :: "r"((p < n_iters) ? 1 : 0), "r"(sg + (unsigned int)(p * 1152)), "l"(g + ((unsigned long long)rowq_ld * (unsigned long long)g_stride + (unsigned long long)(tir * 8)))); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.cg.shared::cta.global [%1], [%2], 16;\n\t" + "}" + :: "r"((p < n_iters) ? 1 : 0), "r"(sg + (unsigned int)(p * 1152 + 512)), "l"(x + ((unsigned long long)rowq_ld * (unsigned long long)x_stride + (unsigned long long)(tir * 8)))); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.ca.shared::cta.global [%1], [%2], 4;\n\t" + "}" + :: "r"((p < n_iters) ? 1 : 0), "r"(sr + (unsigned int)(p * 1152)), "l"(r + rowq_ld)); + asm volatile("cp.async.commit_group;"); + } + int rowx = row_begin + 3 * rstep + rg; + int rowc = row_begin + rg; + int nok = (row_end - row_begin - rg + rstep - 1) / rstep; + if (nok > n_iters) { + nok = n_iters; + } + #pragma unroll 1 + for (int it = 0; it < n_iters; it += 2) { + unsigned int slx = (unsigned int)(it + 3 & 7) * 1152; + int rld = ((rowx) < (last_row) ? (rowx) : (last_row)); + unsigned long long gld = (unsigned long long)rld * (unsigned long long)g_stride + (unsigned long long)(tir * 8); + unsigned long long xld = (unsigned long long)rld * (unsigned long long)x_stride + (unsigned long long)(tir * 8); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.cg.shared::cta.global [%1], [%2], 16;\n\t" + "}" + :: "r"((it + 3 < n_iters) ? 1 : 0), "r"(sg + slx), "l"(g + gld)); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.cg.shared::cta.global [%1], [%2], 16;\n\t" + "}" + :: "r"((it + 3 < n_iters) ? 1 : 0), "r"(sg + slx + 512), "l"(x + xld)); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.ca.shared::cta.global [%1], [%2], 4;\n\t" + "}" + :: "r"((it + 3 < n_iters) ? 1 : 0), "r"(sr + slx), "l"(r + rld)); + asm volatile("cp.async.commit_group;"); + rowx = rowx + rstep; + unsigned int slx_0 = (unsigned int)(it + 3 + 1 & 7) * 1152; + int rld_0 = ((rowx) < (last_row) ? (rowx) : (last_row)); + unsigned long long gld_0 = (unsigned long long)rld_0 * (unsigned long long)g_stride + (unsigned long long)(tir * 8); + unsigned long long xld_0 = (unsigned long long)rld_0 * (unsigned long long)x_stride + (unsigned long long)(tir * 8); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.cg.shared::cta.global [%1], [%2], 16;\n\t" + "}" + :: "r"((it + 3 + 1 < n_iters) ? 1 : 0), "r"(sg + slx_0), "l"(g + gld_0)); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.cg.shared::cta.global [%1], [%2], 16;\n\t" + "}" + :: "r"((it + 3 + 1 < n_iters) ? 1 : 0), "r"(sg + slx_0 + 512), "l"(x + xld_0)); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %0, 0;\n\t" + "@p cp.async.ca.shared::cta.global [%1], [%2], 4;\n\t" + "}" + :: "r"((it + 3 + 1 < n_iters) ? 1 : 0), "r"(sr + slx_0), "l"(r + rld_0)); + asm volatile("cp.async.commit_group;"); + rowx = rowx + rstep; + asm volatile("cp.async.wait_group 3;"); + unsigned int slc = (unsigned int)(it & 7) * 1152; + unsigned int slb = (unsigned int)(it + 1 & 7) * 1152; + __int128_t _smem_raw_3; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_3) : "r"(sg + slc)); + gwA[0] = reinterpret_cast(&_smem_raw_3)[0]; + gwA[1] = reinterpret_cast(&_smem_raw_3)[1]; + gwA[2] = reinterpret_cast(&_smem_raw_3)[2]; + gwA[3] = reinterpret_cast(&_smem_raw_3)[3]; + __int128_t _smem_raw_4; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_4) : "r"(sg + slc + 512)); + xwA[0] = reinterpret_cast(&_smem_raw_4)[0]; + xwA[1] = reinterpret_cast(&_smem_raw_4)[1]; + xwA[2] = reinterpret_cast(&_smem_raw_4)[2]; + xwA[3] = reinterpret_cast(&_smem_raw_4)[3]; + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&rrA[0])) : "r"(sr + slc)); + __int128_t _smem_raw_5; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_5) : "r"(sg + slb)); + gwB[0] = reinterpret_cast(&_smem_raw_5)[0]; + gwB[1] = reinterpret_cast(&_smem_raw_5)[1]; + gwB[2] = reinterpret_cast(&_smem_raw_5)[2]; + gwB[3] = reinterpret_cast(&_smem_raw_5)[3]; + __int128_t _smem_raw_6; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_6) : "r"(sg + slb + 512)); + xwB[0] = reinterpret_cast(&_smem_raw_6)[0]; + xwB[1] = reinterpret_cast(&_smem_raw_6)[1]; + xwB[2] = reinterpret_cast(&_smem_raw_6)[2]; + xwB[3] = reinterpret_cast(&_smem_raw_6)[3]; + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&rrB[0])) : "r"(sr + slb)); + if (it + 2 <= nok) { + float dq = 0.0f; + float gwr[8]; + #pragma unroll + for (int v_1 = 0; v_1 < 1; v_1++) { + unsigned int wvd[4]; + #pragma unroll + for (int e = 0; e < 4; e++) { + unsigned int gwd = gwA[v_1 * 4 + e]; + unsigned int xwd = xwA[v_1 * 4 + e]; + float g0 = __uint_as_float(gwd << 16); + float g1 = __uint_as_float(gwd & 4294901760u); + float x0 = __uint_as_float(xwd << 16); + float x1 = __uint_as_float(xwd & 4294901760u); + float2 wq = make_float2(wc[v_1 * 8 + 2 * e], wc[v_1 * 8 + 2 * e + 1]); + float2 _f32x2_mul_a_7 = make_float2(g0, g1); + float2 gwq; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&gwq) : "l"(*(const unsigned long long*)&_f32x2_mul_a_7), "l"(*(const unsigned long long*)&wq)); + gwr[v_1 * 8 + 2 * e] = gwq.x; + gwr[v_1 * 8 + 2 * e + 1] = gwq.y; + dq = __fmaf_rn(gwq.x, x0, dq); + dq = __fmaf_rn(gwq.y, x1, dq); + } + } + float dots[1]; + float _warp_reduce_8 = dq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_8 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_8, offset); + dots[0] = _warp_reduce_8; + float rq = rrA[0]; + float cq = dots[0] * 0.00390625f; + float kx = rq * rq * rq * cq; + float nkx = -kx; + unsigned long long dx_base = (unsigned long long)rowc * 256; + float2 rr = make_float2(rq, rq); + float2 rrn = make_float2(-rq, -rq); + float2 nn = make_float2(nkx, nkx); + #pragma unroll + for (int v_2 = 0; v_2 < 1; v_2++) { + float out[8]; + unsigned int wv[4]; + #pragma unroll + for (int e_1 = 0; e_1 < 4; e_1++) { + unsigned int gwd_1 = gwA[v_2 * 4 + e_1]; + unsigned int xwd_1 = xwA[v_2 * 4 + e_1]; + float g0_1 = __uint_as_float(gwd_1 << 16); + float g1_1 = __uint_as_float(gwd_1 & 4294901760u); + float x0_1 = __uint_as_float(xwd_1 << 16); + float x1_1 = __uint_as_float(xwd_1 & 4294901760u); + float2 _f32x2_mul_a_9 = make_float2(g0_1, g1_1); + float2 _f32x2_mul_b_10 = make_float2(x0_1, x1_1); + float2 gx; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&gx) : "l"(*(const unsigned long long*)&_f32x2_mul_a_9), "l"(*(const unsigned long long*)&_f32x2_mul_b_10)); + float2 pn; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&pn) : "l"(*(const unsigned long long*)&gx), "l"(*(const unsigned long long*)&rrn)); + float2 pe = fma_f32x2_rn_noftz(gx, rr, pn); + float2 a2 = make_float2(acc[v_2 * 8 + 2 * e_1], acc[v_2 * 8 + 2 * e_1 + 1]); + float2 s2 = sub_f32x2_noftz(a2, pn); + float2 z2 = sub_f32x2_noftz(s2, a2); + float2 t2 = sub_f32x2_noftz(pe, add_f32x2_noftz(pn, z2)); + float2 c2 = add_f32x2_noftz(make_float2(comp[v_2 * 8 + 2 * e_1], comp[v_2 * 8 + 2 * e_1 + 1]), t2); + acc[v_2 * 8 + 2 * e_1] = s2.x; + acc[v_2 * 8 + 2 * e_1 + 1] = s2.y; + comp[v_2 * 8 + 2 * e_1] = c2.x; + comp[v_2 * 8 + 2 * e_1 + 1] = c2.y; + float2 _f32x2_mul_b_11 = make_float2(gwr[v_2 * 8 + 2 * e_1], gwr[v_2 * 8 + 2 * e_1 + 1]); + float2 q2; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&q2) : "l"(*(const unsigned long long*)&rr), "l"(*(const unsigned long long*)&_f32x2_mul_b_11)); + float2 o2 = fma_f32x2_rn_noftz(nn, make_float2(x0_1, x1_1), q2); + out[2 * e_1] = o2.x; + out[2 * e_1 + 1] = o2.y; + } + int col = (v_2 * 32 + tir) * 8; + unsigned int packed[4]; + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out[_lp*2 + 0], out[_lp*2+1 + 0])); + packed[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out[_lp*2 + 2], out[_lp*2+1 + 2])); + packed[_lp + 1] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out[_lp*2 + 4], out[_lp*2+1 + 4])); + packed[_lp + 2] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out[_lp*2 + 6], out[_lp*2+1 + 6])); + packed[_lp + 3] = *(uint32_t*)&_bf2; + } + reinterpret_cast(dx + (dx_base + (unsigned long long)col))[0] = reinterpret_cast(packed)[0]; + } + float dq_1 = 0.0f; + float gwr_0[8]; + #pragma unroll + for (int v_3 = 0; v_3 < 1; v_3++) { + unsigned int wvd_1[4]; + #pragma unroll + for (int e_2 = 0; e_2 < 4; e_2++) { + unsigned int gwd_2 = gwB[v_3 * 4 + e_2]; + unsigned int xwd_2 = xwB[v_3 * 4 + e_2]; + float g0_2 = __uint_as_float(gwd_2 << 16); + float g1_2 = __uint_as_float(gwd_2 & 4294901760u); + float x0_2 = __uint_as_float(xwd_2 << 16); + float x1_2 = __uint_as_float(xwd_2 & 4294901760u); + float2 wq_1 = make_float2(wc[v_3 * 8 + 2 * e_2], wc[v_3 * 8 + 2 * e_2 + 1]); + float2 _f32x2_mul_a_12 = make_float2(g0_2, g1_2); + float2 gwq_1; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&gwq_1) : "l"(*(const unsigned long long*)&_f32x2_mul_a_12), "l"(*(const unsigned long long*)&wq_1)); + gwr_0[v_3 * 8 + 2 * e_2] = gwq_1.x; + gwr_0[v_3 * 8 + 2 * e_2 + 1] = gwq_1.y; + dq_1 = __fmaf_rn(gwq_1.x, x0_2, dq_1); + dq_1 = __fmaf_rn(gwq_1.y, x1_2, dq_1); + } + } + float dots_0[1]; + float _warp_reduce_13 = dq_1; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_13 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_13, offset); + dots_0[0] = _warp_reduce_13; + float rq_0 = rrB[0]; + float cq_0 = dots_0[0] * 0.00390625f; + float kx_0 = rq_0 * rq_0 * rq_0 * cq_0; + float nkx_0 = -kx_0; + unsigned long long dx_base_0 = (unsigned long long)(rowc + rstep) * 256; + float2 rr_0 = make_float2(rq_0, rq_0); + float2 rrn_0 = make_float2(-rq_0, -rq_0); + float2 nn_0 = make_float2(nkx_0, nkx_0); + #pragma unroll + for (int v_4 = 0; v_4 < 1; v_4++) { + float out_1[8]; + unsigned int wv_1[4]; + #pragma unroll + for (int e_3 = 0; e_3 < 4; e_3++) { + unsigned int gwd_3 = gwB[v_4 * 4 + e_3]; + unsigned int xwd_3 = xwB[v_4 * 4 + e_3]; + float g0_3 = __uint_as_float(gwd_3 << 16); + float g1_3 = __uint_as_float(gwd_3 & 4294901760u); + float x0_3 = __uint_as_float(xwd_3 << 16); + float x1_3 = __uint_as_float(xwd_3 & 4294901760u); + float2 _f32x2_mul_a_14 = make_float2(g0_3, g1_3); + float2 _f32x2_mul_b_15 = make_float2(x0_3, x1_3); + float2 gx_1; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&gx_1) : "l"(*(const unsigned long long*)&_f32x2_mul_a_14), "l"(*(const unsigned long long*)&_f32x2_mul_b_15)); + float2 pn_1; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&pn_1) : "l"(*(const unsigned long long*)&gx_1), "l"(*(const unsigned long long*)&rrn_0)); + float2 pe_1 = fma_f32x2_rn_noftz(gx_1, rr_0, pn_1); + float2 a2_1 = make_float2(acc[v_4 * 8 + 2 * e_3], acc[v_4 * 8 + 2 * e_3 + 1]); + float2 s2_1 = sub_f32x2_noftz(a2_1, pn_1); + float2 z2_1 = sub_f32x2_noftz(s2_1, a2_1); + float2 t2_1 = sub_f32x2_noftz(pe_1, add_f32x2_noftz(pn_1, z2_1)); + float2 c2_1 = add_f32x2_noftz(make_float2(comp[v_4 * 8 + 2 * e_3], comp[v_4 * 8 + 2 * e_3 + 1]), t2_1); + acc[v_4 * 8 + 2 * e_3] = s2_1.x; + acc[v_4 * 8 + 2 * e_3 + 1] = s2_1.y; + comp[v_4 * 8 + 2 * e_3] = c2_1.x; + comp[v_4 * 8 + 2 * e_3 + 1] = c2_1.y; + float2 _f32x2_mul_b_16 = make_float2(gwr_0[v_4 * 8 + 2 * e_3], gwr_0[v_4 * 8 + 2 * e_3 + 1]); + float2 q2_1; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&q2_1) : "l"(*(const unsigned long long*)&rr_0), "l"(*(const unsigned long long*)&_f32x2_mul_b_16)); + float2 o2_1 = fma_f32x2_rn_noftz(nn_0, make_float2(x0_3, x1_3), q2_1); + out_1[2 * e_3] = o2_1.x; + out_1[2 * e_3 + 1] = o2_1.y; + } + int col_1 = (v_4 * 32 + tir) * 8; + unsigned int packed_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_1[_lp*2 + 0], out_1[_lp*2+1 + 0])); + packed_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_1[_lp*2 + 2], out_1[_lp*2+1 + 2])); + packed_1[_lp + 1] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_1[_lp*2 + 4], out_1[_lp*2+1 + 4])); + packed_1[_lp + 2] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_1[_lp*2 + 6], out_1[_lp*2+1 + 6])); + packed_1[_lp + 3] = *(uint32_t*)&_bf2; + } + reinterpret_cast(dx + (dx_base_0 + (unsigned long long)col_1))[0] = reinterpret_cast(packed_1)[0]; + } + } else if (it < nok) { + float dq_2 = 0.0f; + float gwr_1[8]; + #pragma unroll + for (int v_5 = 0; v_5 < 1; v_5++) { + unsigned int wvd_2[4]; + #pragma unroll + for (int e_4 = 0; e_4 < 4; e_4++) { + unsigned int gwd_4 = gwA[v_5 * 4 + e_4]; + unsigned int xwd_4 = xwA[v_5 * 4 + e_4]; + float g0_4 = __uint_as_float(gwd_4 << 16); + float g1_4 = __uint_as_float(gwd_4 & 4294901760u); + float x0_4 = __uint_as_float(xwd_4 << 16); + float x1_4 = __uint_as_float(xwd_4 & 4294901760u); + float2 wq_2 = make_float2(wc[v_5 * 8 + 2 * e_4], wc[v_5 * 8 + 2 * e_4 + 1]); + float2 _f32x2_mul_a_17 = make_float2(g0_4, g1_4); + float2 gwq_2; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&gwq_2) : "l"(*(const unsigned long long*)&_f32x2_mul_a_17), "l"(*(const unsigned long long*)&wq_2)); + gwr_1[v_5 * 8 + 2 * e_4] = gwq_2.x; + gwr_1[v_5 * 8 + 2 * e_4 + 1] = gwq_2.y; + dq_2 = __fmaf_rn(gwq_2.x, x0_4, dq_2); + dq_2 = __fmaf_rn(gwq_2.y, x1_4, dq_2); + } + } + float dots_1[1]; + float _warp_reduce_18 = dq_2; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_18 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_18, offset); + dots_1[0] = _warp_reduce_18; + float rq_1 = rrA[0]; + float cq_1 = dots_1[0] * 0.00390625f; + float kx_1 = rq_1 * rq_1 * rq_1 * cq_1; + float nkx_1 = -kx_1; + unsigned long long dx_base_1 = (unsigned long long)rowc * 256; + float2 rr_1 = make_float2(rq_1, rq_1); + float2 rrn_1 = make_float2(-rq_1, -rq_1); + float2 nn_1 = make_float2(nkx_1, nkx_1); + #pragma unroll + for (int v_6 = 0; v_6 < 1; v_6++) { + float out_2[8]; + unsigned int wv_2[4]; + #pragma unroll + for (int e_5 = 0; e_5 < 4; e_5++) { + unsigned int gwd_5 = gwA[v_6 * 4 + e_5]; + unsigned int xwd_5 = xwA[v_6 * 4 + e_5]; + float g0_5 = __uint_as_float(gwd_5 << 16); + float g1_5 = __uint_as_float(gwd_5 & 4294901760u); + float x0_5 = __uint_as_float(xwd_5 << 16); + float x1_5 = __uint_as_float(xwd_5 & 4294901760u); + float2 _f32x2_mul_a_19 = make_float2(g0_5, g1_5); + float2 _f32x2_mul_b_20 = make_float2(x0_5, x1_5); + float2 gx_2; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&gx_2) : "l"(*(const unsigned long long*)&_f32x2_mul_a_19), "l"(*(const unsigned long long*)&_f32x2_mul_b_20)); + float2 pn_2; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&pn_2) : "l"(*(const unsigned long long*)&gx_2), "l"(*(const unsigned long long*)&rrn_1)); + float2 pe_2 = fma_f32x2_rn_noftz(gx_2, rr_1, pn_2); + float2 a2_2 = make_float2(acc[v_6 * 8 + 2 * e_5], acc[v_6 * 8 + 2 * e_5 + 1]); + float2 s2_2 = sub_f32x2_noftz(a2_2, pn_2); + float2 z2_2 = sub_f32x2_noftz(s2_2, a2_2); + float2 t2_2 = sub_f32x2_noftz(pe_2, add_f32x2_noftz(pn_2, z2_2)); + float2 c2_2 = add_f32x2_noftz(make_float2(comp[v_6 * 8 + 2 * e_5], comp[v_6 * 8 + 2 * e_5 + 1]), t2_2); + acc[v_6 * 8 + 2 * e_5] = s2_2.x; + acc[v_6 * 8 + 2 * e_5 + 1] = s2_2.y; + comp[v_6 * 8 + 2 * e_5] = c2_2.x; + comp[v_6 * 8 + 2 * e_5 + 1] = c2_2.y; + float2 _f32x2_mul_b_21 = make_float2(gwr_1[v_6 * 8 + 2 * e_5], gwr_1[v_6 * 8 + 2 * e_5 + 1]); + float2 q2_2; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&q2_2) : "l"(*(const unsigned long long*)&rr_1), "l"(*(const unsigned long long*)&_f32x2_mul_b_21)); + float2 o2_2 = fma_f32x2_rn_noftz(nn_1, make_float2(x0_5, x1_5), q2_2); + out_2[2 * e_5] = o2_2.x; + out_2[2 * e_5 + 1] = o2_2.y; + } + int col_2 = (v_6 * 32 + tir) * 8; + unsigned int packed_2[4]; + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_2[_lp*2 + 0], out_2[_lp*2+1 + 0])); + packed_2[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_2[_lp*2 + 2], out_2[_lp*2+1 + 2])); + packed_2[_lp + 1] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_2[_lp*2 + 4], out_2[_lp*2+1 + 4])); + packed_2[_lp + 2] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(out_2[_lp*2 + 6], out_2[_lp*2+1 + 6])); + packed_2[_lp + 3] = *(uint32_t*)&_bf2; + } + reinterpret_cast(dx + (dx_base_1 + (unsigned long long)col_2))[0] = reinterpret_cast(packed_2)[0]; + } + } + rowc = rowc + 2 * rstep; + } + asm volatile("cp.async.wait_group 0;"); + } + unsigned long long pbase = (unsigned long long)c * 256; + #pragma unroll + for (int v_7 = 0; v_7 < 1; v_7++) { + int fcol = (v_7 * 32 + tir) * 8; + #pragma unroll + for (int e_6 = 0; e_6 < 8; e_6++) { + row_fold[rg * 256 + fcol + e_6] = acc[v_7 * 8 + e_6]; + } + } + __syncthreads(); + float fold_acc[2]; + #pragma unroll + for (int e_7 = 0; e_7 < 2; e_7++) { + fold_acc[e_7] = 0.0f; + } + float fold_c[2]; + #pragma unroll + for (int e_8 = 0; e_8 < 2; e_8++) { + fold_c[e_8] = 0.0f; + } + #pragma unroll + for (int k_1 = 0; k_1 < 4; k_1++) { + #pragma unroll + for (int e_9 = 0; e_9 < 2; e_9++) { + float d2 = row_fold[k_1 * 256 + tid * 2 + e_9]; + float fs = fold_acc[e_9] + d2; + float fb = fs - fold_acc[e_9]; + fold_c[e_9] = fold_c[e_9] + (fold_acc[e_9] - (fs - fb) + (d2 - fb)); + fold_acc[e_9] = fs; + } + } + __syncthreads(); + #pragma unroll + for (int e5 = 0; e5 < 8; e5++) { + row_fold[rg * 256 + tir * 8 + e5] = comp[e5]; + } + __syncthreads(); + #pragma unroll + for (int k5 = 0; k5 < 4; k5++) { + #pragma unroll + for (int e_10 = 0; e_10 < 2; e_10++) { + float d5 = row_fold[k5 * 256 + tid * 2 + e_10]; + fold_c[e_10] = fold_c[e_10] + d5; + } + } + #pragma unroll + for (int e_11 = 0; e_11 < 2; e_11++) { + fold_acc[e_11] = fold_acc[e_11] + fold_c[e_11]; + } + int pcol = tid * 2; + { + float2 _v2 = make_float2(fold_acc[0 + 0], fold_acc[0 + 1]); + *reinterpret_cast(partial + pbase + (unsigned long long)pcol) = _v2; + } + if (tid == 0) { + flags[claim_it & 1] = (int)nxt; + } + __syncthreads(); + c = flags[claim_it & 1]; + } + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + if (tid == 0) { + if (c == n_chunks + (int)gridDim.x - 1) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } +} + +} // extern "C" diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_binding.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_binding.cu new file mode 100644 index 00000000000..8a566ba99eb --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_binding.cu @@ -0,0 +1,146 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05, for the device source that defines it. Do not edit. +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, float* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0); + +namespace cake_binding_kernel_cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05 { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckRank(const TensorView& t, const char* name, int rank) { + TVM_FFI_CHECK(t.ndim() >= rank, ValueError) + << KERNEL_NAME << ": '" << name << "' needs at least " << rank << " dims, got " << t.ndim(); +} + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +void Run(TensorView a_g, TensorView a_x, TensorView a_w, TensorView a_r, TensorView a_dx, TensorView a_dw, TensorView a_partial, TensorView a_counters, TensorView a_g_h, int64_t a_T, int64_t a_n_chunks, int64_t a_rows_per_chunk, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_g.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_g, "g"); + check_dtype(a_g, DLDataType{kDLBfloat, 16, 1}, "g"); + check_cuda_tensor(a_x, "x"); + check_dtype(a_x, DLDataType{kDLBfloat, 16, 1}, "x"); + check_cuda_tensor(a_w, "w"); + check_dtype(a_w, DLDataType{kDLFloat, 32, 1}, "w"); + check_cuda_tensor(a_r, "r"); + check_dtype(a_r, DLDataType{kDLFloat, 32, 1}, "r"); + check_cuda_tensor(a_dx, "dx"); + check_dtype(a_dx, DLDataType{kDLBfloat, 16, 1}, "dx"); + check_cuda_tensor(a_dw, "dw"); + check_dtype(a_dw, DLDataType{kDLFloat, 32, 1}, "dw"); + check_cuda_tensor(a_partial, "partial"); + check_dtype(a_partial, DLDataType{kDLFloat, 32, 1}, "partial"); + check_cuda_tensor(a_counters, "counters"); + check_dtype(a_counters, DLDataType{kDLUInt, 32, 1}, "counters"); + check_cuda_tensor(a_g_h, "g_h"); + check_dtype(a_g_h, DLDataType{kDLBfloat, 16, 1}, "g_h"); + TVM_FFI_CHECK(a_T >= -2147483648LL && a_T <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'T' = " << a_T + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_n_chunks >= -2147483648LL && a_n_chunks <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'n_chunks' = " << a_n_chunks + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_rows_per_chunk >= -2147483648LL && a_rows_per_chunk <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'rows_per_chunk' = " << a_rows_per_chunk + << " is outside the i32 range [-2147483648, 2147483647]"; + CheckRank(a_g, "g", 1); + CheckUnitInnerStride(a_g, "g"); + CheckRank(a_x, "x", 1); + CheckUnitInnerStride(a_x, "x"); + check_same_device(a_x, a_g, "x", "g"); + check_contiguous(a_w, "w"); + check_same_device(a_w, a_g, "w", "g"); + check_contiguous(a_r, "r"); + check_same_device(a_r, a_g, "r", "g"); + check_contiguous(a_dx, "dx"); + check_same_device(a_dx, a_g, "dx", "g"); + check_contiguous(a_dw, "dw"); + check_same_device(a_dw, a_g, "dw", "g"); + check_contiguous(a_partial, "partial"); + check_same_device(a_partial, a_g, "partial", "g"); + check_contiguous(a_counters, "counters"); + check_same_device(a_counters, a_g, "counters", "g"); + CheckRank(a_g_h, "g_h", 1); + CheckUnitInnerStride(a_g_h, "g_h"); + check_same_device(a_g_h, a_g, "g_h", "g"); + TVM_FFI_CHECK((a_n_chunks >= 1LL), ValueError) + << KERNEL_NAME << ": requires (n_chunks >= 1)"; + TVM_FFI_CHECK((a_rows_per_chunk >= 1LL), ValueError) + << KERNEL_NAME << ": requires (rows_per_chunk >= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05), cudaFuncAttributeMaxDynamicSharedMemorySize, 99840), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + void* v_g = DataPtr(a_g); + void* v_x = DataPtr(a_x); + void* v_w = DataPtr(a_w); + void* v_r = DataPtr(a_r); + void* v_dx = DataPtr(a_dx); + void* v_dw = DataPtr(a_dw); + void* v_partial = DataPtr(a_partial); + void* v_counters = DataPtr(a_counters); + void* v_g_h = DataPtr(a_g_h); + int32_t v_T = static_cast(a_T); + int32_t v_n_chunks = static_cast(a_n_chunks); + int32_t v_rows_per_chunk = static_cast(a_rows_per_chunk); + int64_t v_g_stride_p0 = a_g.stride(0); + int64_t v_x_stride_p0 = a_x.stride(0); + int64_t v_g_h_stride_p0 = a_g_h.stride(0); + void* kernel_args[] = {&v_g, &v_x, &v_w, &v_r, &v_dx, &v_dw, &v_partial, &v_counters, &v_g_h, &v_T, &v_n_chunks, &v_rows_per_chunk, &v_g_stride_p0, &v_x_stride_p0, &v_g_h_stride_p0}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(384, 1, 1); + config.dynamicSmemBytes = 99840; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_kernel.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_kernel.cu new file mode 100644 index 00000000000..af0b3baaa15 --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05_kernel.cu @@ -0,0 +1,819 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define SMEM_STAGE_BUF_OFF 1024 +#define SMEM_STAGE_BUF_STAGE_BYTES 73728 +#define SMEM_STAGE_BUF_STRIDE 73728 +#define SMEM_ROW_FOLD_OFF 1024 +#define SMEM_ROW_FOLD_STAGE_BYTES 24576 +#define SMEM_ROW_FOLD_STRIDE 24576 +#define SMEM_RED_FOLD_OFF 25600 +#define SMEM_RED_FOLD_STAGE_BYTES 6144 +#define SMEM_RED_FOLD_STRIDE 6144 +#define SMEM_DOT_SLOTS_OFF 74752 +#define SMEM_DOT_SLOTS_STAGE_BYTES 96 +#define SMEM_DOT_SLOTS_STRIDE 96 +#define SMEM_PART_SM_OFF 74880 +#define SMEM_PART_SM_STAGE_BYTES 16 +#define SMEM_PART_SM_STRIDE 16 +#define SMEM_DESC_OFF 99584 +#define SMEM_DESC_STAGE_BYTES 60 +#define SMEM_DESC_STRIDE 60 +#define SMEM_RTAB_OFF 99712 +#define SMEM_RTAB_STAGE_BYTES 128 +#define SMEM_RTAB_STRIDE 128 +#define SMEM_FOLD_SM_OFF 75008 +#define SMEM_FOLD_SM_STAGE_BYTES 24576 +#define SMEM_FOLD_SM_STRIDE 24576 +#define SMEM_TOTAL 99840 +#define THREADS 384 +#define LAUNCH_MIN_BLOCKS 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_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"); +} + + + +extern "C" { + +__global__ __launch_bounds__(384, LAUNCH_MIN_BLOCKS) void +kernel_cake_rmsnorm_det_train_ca6211f3ea2ce8c86a05(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, float* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define bars_addr (mbar_base + 0) + #define dbar_addr (mbar_base + 24) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + uint8_t* stage_buf = reinterpret_cast(smem_raw + 1024); + const int stage_buf_addr = smem + 1024; + float* row_fold = reinterpret_cast(smem_raw + 1024); + const int row_fold_addr = smem + 1024; + float* red_fold = reinterpret_cast(smem_raw + 25600); + const int red_fold_addr = smem + 25600; + float* dot_slots = reinterpret_cast(smem_raw + 74752); + const int dot_slots_addr = smem + 74752; + float* part_sm = reinterpret_cast(smem_raw + 74880); + const int part_sm_addr = smem + 74880; + int* desc = reinterpret_cast(smem_raw + 99584); + const int desc_addr = smem + 99584; + int* rtab = reinterpret_cast(smem_raw + 99712); + const int rtab_addr = smem + 99712; + float* fold_sm = reinterpret_cast(smem_raw + 75008); + const int fold_sm_addr = smem + 75008; + + // Mbarrier init (2 pipeline groups, 0 ordered-sequence groups, 8 barriers) + // Mbarriers at smem_raw[0..64) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // bars: 3 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + mbarrier_init(smem + 8, 1); + mbarrier_init(smem + 16, 1); + // dbar: 5 barriers, init_count=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); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncthreads(); + + // === Task calls (declaration order) === + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + long long g_stride = g_stride_p0; + long long x_stride = x_stride_p0; + long long gh_stride = g_h_stride_p0; + int c = bid; + int rg = tid / 96; + int tir = tid - rg * 96; + unsigned int claim[1]; + if (tid == 0) { + unsigned int _atomic_fetch_add_0 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_0; + int S = (T + 3) / 4; + int _min_1 = ((rows_per_chunk) < (S) ? (rows_per_chunk) : (S)); + int sl = ((_min_1) > (1) ? (_min_1) : (1)); + int share = (S + sl - 1) / sl; + int tail = (((share + 11) / 12) > (1) ? ((share + 11) / 12) : (1)); + int size = (share + 1) / 2; + int rem = S; + int base = 0; + int first = 0; + int nr = 0; + int walk = 1; + while (walk == 1) { + walk = 0; + if (nr < 5) { + if (size > tail) { + if (rem >= sl * size) { + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = size; + first += sl; + base += sl * size; + rem -= sl * size; + size = (size + 1) / 2; + nr += 1; + walk = 1; + } + } + } + } + int q = rem / sl; + int m = rem - q * sl; + if (m > 0) { + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = q + 1; + first += m; + base += m * (q + 1); + nr += 1; + } + if (q > 0) { + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = q; + first += sl - m; + base += (sl - m) * q; + nr += 1; + } + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = 0; + rtab[31] = nr + 1; + } + int ring = (int)stage_buf_addr; + int my_g = ring + rg * 3072 + tir * 16; + float wf[16]; + #pragma unroll + for (int v = 0; v < 2; v++) { + int wcol = (v * 96 + tir) * 8; + { + unsigned _ldv8_2_0; + unsigned _ldv8_2_1; + unsigned _ldv8_2_2; + unsigned _ldv8_2_3; + unsigned _ldv8_2_4; + unsigned _ldv8_2_5; + unsigned _ldv8_2_6; + unsigned _ldv8_2_7; + asm volatile( + "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_2_0), "=r"(_ldv8_2_1), "=r"(_ldv8_2_2), "=r"(_ldv8_2_3), "=r"(_ldv8_2_4), "=r"(_ldv8_2_5), "=r"(_ldv8_2_6), "=r"(_ldv8_2_7) : "l"((const void*)(w + ((unsigned long long)wcol))) : "memory"); + wf[v * 8 + 0] = __uint_as_float(_ldv8_2_0); + wf[v * 8 + 1] = __uint_as_float(_ldv8_2_1); + wf[v * 8 + 2] = __uint_as_float(_ldv8_2_2); + wf[v * 8 + 3] = __uint_as_float(_ldv8_2_3); + wf[v * 8 + 4] = __uint_as_float(_ldv8_2_4); + wf[v * 8 + 5] = __uint_as_float(_ldv8_2_5); + wf[v * 8 + 6] = __uint_as_float(_ldv8_2_6); + wf[v * 8 + 7] = __uint_as_float(_ldv8_2_7); + } + } + int rng[2]; + int nd[3]; + if (tid == 0) { + int cid0 = (int)claim[0]; + int nr_1 = rtab[31]; + int r_0 = 0; + int walk_1 = 1; + while (walk_1 == 1) { + walk_1 = 0; + if (nr_1 > r_0 + 2) { + if (cid0 >= rtab[(r_0 + 1) * 3]) { + r_0 += 1; + walk_1 = 1; + } + } + } + if (cid0 >= rtab[(nr_1 - 1) * 3]) { + rng[0] = T; + rng[1] = T; + } else { + int first_1 = rtab[r_0 * 3]; + int base_1 = rtab[r_0 * 3 + 1]; + int stages = rtab[r_0 * 3 + 2]; + rng[0] = (base_1 + (cid0 - first_1) * stages) * 4; + int _min_3 = ((rng[0] + stages * 4) < (T) ? (rng[0] + stages * 4) : (T)); + rng[1] = _min_3; + } + int entry = 0; + int word = entry * 3; + desc[word] = cid0; + desc[word + 1] = rng[0]; + desc[word + 2] = rng[1]; + mbarrier_arrive(dbar_addr + (entry) * 8); + if (rng[0] < T) { + unsigned int _atomic_fetch_add_4 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_4; + } else { + int nr_0 = rtab[31]; + int last = rtab[(nr_0 - 1) * 3] + num_bids - 1; + if (T == 0) { + last = 0; + } + if (cid0 == last) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } + } + int entry_1 = 0; + mbarrier_wait(dbar_addr + (entry_1) * 8, 0); + int word_1 = entry_1 * 3; + int w0 = desc[word_1]; + int w1 = desc[word_1 + 1]; + int w2 = desc[word_1 + 2]; + int _warp_redux_5; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_5) : "r"(w0)); + nd[0] = _warp_redux_5; + int _warp_redux_6; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_6) : "r"(w1)); + nd[1] = _warp_redux_6; + int _warp_redux_7; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_7) : "r"(w2)); + nd[2] = _warp_redux_7; + int c0 = nd[0]; + int cb = nd[1]; + int ce = nd[2]; + int ps[3]; + ps[0] = nd[1]; + ps[1] = nd[2]; + int desc_n = 0; + #pragma unroll + for (int s = 0; s < 3; s++) { + ps[2] = 0; + if (ps[0] < ps[1]) { + int first_2 = ps[0]; + int rows_k = ((ps[1] - first_2) < (4) ? (ps[1] - first_2) : (4)); + int sbase = ring + s * 24576; + if (tid == 0) { + mbarrier_arrive_expect_tx(bars_addr + (s) * 8, rows_k * 6144); + } + if (tir == 0) { + if (rg < rows_k) { + int row_g = first_2 + rg; + long long g_srcg = (long long)row_g * g_stride; + long long x_srcg = (long long)row_g * x_stride; + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase + rg * 3072), "l"(reinterpret_cast(g + g_srcg)), "r"((unsigned)(3072)), "r"(bars_addr + (s) * 8), + "l"(0x12F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase + (rg + 4) * 3072), "l"(reinterpret_cast(x + x_srcg)), "r"((unsigned)(3072)), "r"(bars_addr + (s) * 8), + "l"(0x12F0000000000000ULL) : "memory"); + } + } + ps[0] = ps[0] + 4; + ps[2] = (int)(ps[0] >= ps[1]); + if (ps[0] >= ps[1]) { + if (tid == 0) { + int cid = (int)claim[0]; + int nr_2 = rtab[31]; + int r_0_1 = 0; + int walk_2 = 1; + while (walk_2 == 1) { + walk_2 = 0; + if (nr_2 > r_0_1 + 2) { + if (cid >= rtab[(r_0_1 + 1) * 3]) { + r_0_1 += 1; + walk_2 = 1; + } + } + } + if (cid >= rtab[(nr_2 - 1) * 3]) { + rng[0] = T; + rng[1] = T; + } else { + int first_0 = rtab[r_0_1 * 3]; + int base_2 = rtab[r_0_1 * 3 + 1]; + int stages_1 = rtab[r_0_1 * 3 + 2]; + rng[0] = (base_2 + (cid - first_0) * stages_1) * 4; + int _min_8 = ((rng[0] + stages_1 * 4) < (T) ? (rng[0] + stages_1 * 4) : (T)); + rng[1] = _min_8; + } + int entry_1_1 = (desc_n + 1) % 5; + int word_2 = entry_1_1 * 3; + desc[word_2] = cid; + desc[word_2 + 1] = rng[0]; + desc[word_2 + 2] = rng[1]; + mbarrier_arrive(dbar_addr + (entry_1_1) * 8); + if (rng[0] < T) { + unsigned int _atomic_fetch_add_9 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_9; + } else { + int nr_0_1 = rtab[31]; + int last_1 = rtab[(nr_0_1 - 1) * 3] + num_bids - 1; + if (T == 0) { + last_1 = 0; + } + if (cid == last_1) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } + } + int entry_0 = (desc_n + 1) % 5; + int _parity_dbar = (int)((unsigned int)(desc_n + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_0) * 8, _parity_dbar); + int word_1_1 = entry_0 * 3; + int w0_2 = desc[word_1_1]; + int w1_3 = desc[word_1_1 + 1]; + int w2_4 = desc[word_1_1 + 2]; + int _warp_redux_10; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_10) : "r"(w0_2)); + nd[0] = _warp_redux_10; + int _warp_redux_11; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_11) : "r"(w1_3)); + nd[1] = _warp_redux_11; + int _warp_redux_12; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_12) : "r"(w2_4)); + nd[2] = _warp_redux_12; + ps[0] = nd[1]; + ps[1] = nd[2]; + } + } + int desc_step = ps[2]; + desc_n += desc_step; + } + float acc[16]; + #pragma unroll + for (int j = 0; j < 16; j++) { + acc[j] = 0.0f; + } + float cmp[16]; + #pragma unroll + for (int jc = 0; jc < 16; jc++) { + cmp[jc] = 0.0f; + } + float hi[16]; + #pragma unroll + for (int jh = 0; jh < 16; jh++) { + hi[jh] = 0.0f; + } + int n_cons = 0; + int g_step = 0; + float r_cur = 0.0f; + if (cb < ce) { + int row0 = ((cb + rg) < (ce - 1) ? (cb + rg) : (ce - 1)); + r_cur = r[row0]; + } + unsigned int gw[8]; + unsigned int xw[8]; + float out[16]; + if (cb >= ce) { + if (c0 < n_chunks) { + unsigned long long pbase0 = (unsigned long long)c0 * 1536; + float fold0[4]; + #pragma unroll + for (int e0 = 0; e0 < 4; e0++) { + fold0[e0] = 0.0f; + } + int pcol0 = tid * 4; + { + float4 _v4 = make_float4(fold0[0 + 0], fold0[0 + 1], fold0[0 + 2], fold0[0 + 3]); + *reinterpret_cast(partial + pbase0 + (unsigned long long)pcol0) = _v4; + } + } + } + while (cb < ce) { + int n_steps_c = (ce - cb + 3) / 4; + #pragma unroll 1 + for (int k = 0; k < n_steps_c; k++) { + int slot = g_step % 3; + int par = (g_step & 1) * 12; + int row = cb + k * 4 + rg; + int row_n = ((row + 4) < (ce - 1) ? (row + 4) : (ce - 1)); + if (n_steps_c <= k + 1) { + int entry_0_1 = (n_cons + 1) % 5; + int _parity_dbar_1 = (int)((unsigned int)(n_cons + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_0_1) * 8, _parity_dbar_1); + int word_0 = entry_0_1 * 3; + int w0_0 = desc[word_0]; + int w1_0 = desc[word_0 + 1]; + int w2_0 = desc[word_0 + 2]; + int _warp_redux_13; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_13) : "r"(w0_0)); + nd[0] = _warp_redux_13; + int _warp_redux_14; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_14) : "r"(w1_0)); + nd[1] = _warp_redux_14; + int _warp_redux_15; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_15) : "r"(w2_0)); + nd[2] = _warp_redux_15; + row_n = cb; + if (nd[1] < T) { + int _min_16 = ((nd[1] + rg) < (nd[2] - 1) ? (nd[1] + rg) : (nd[2] - 1)); + row_n = _min_16; + } + } + float r_nxt = r[row_n]; + int _parity_bars = (int)((unsigned int)g_step / 3 % 2); + mbarrier_wait(bars_addr + (slot) * 8, _parity_bars); + int sg = my_g + slot * 24576; + #pragma unroll + for (int v_1 = 0; v_1 < 2; v_1++) { + __int128_t _smem_raw_17; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_17) : "r"(sg + v_1 * 1536)); + gw[v_1 * 4] = reinterpret_cast(&_smem_raw_17)[0]; + gw[(v_1 * 4) + 1] = reinterpret_cast(&_smem_raw_17)[1]; + gw[(v_1 * 4) + 2] = reinterpret_cast(&_smem_raw_17)[2]; + gw[(v_1 * 4) + 3] = reinterpret_cast(&_smem_raw_17)[3]; + __int128_t _smem_raw_18; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_18) : "r"(sg + (12288 + v_1 * 1536))); + xw[v_1 * 4] = reinterpret_cast(&_smem_raw_18)[0]; + xw[(v_1 * 4) + 1] = reinterpret_cast(&_smem_raw_18)[1]; + xw[(v_1 * 4) + 2] = reinterpret_cast(&_smem_raw_18)[2]; + xw[(v_1 * 4) + 3] = reinterpret_cast(&_smem_raw_18)[3]; + } + float dq = 0.0f; + #pragma unroll + for (int v_2 = 0; v_2 < 2; v_2++) { + #pragma unroll + for (int e = 0; e < 4; e++) { + unsigned int gwd = gw[v_2 * 4 + e]; + unsigned int xwd = xw[v_2 * 4 + e]; + float g0 = __uint_as_float(gwd << 16); + float g1 = __uint_as_float(gwd & 4294901760u); + float x0 = __uint_as_float(xwd << 16); + float x1 = __uint_as_float(xwd & 4294901760u); + dq = __fmaf_rn(g0 * wf[v_2 * 8 + 2 * e], x0, dq); + dq = __fmaf_rn(g1 * wf[v_2 * 8 + 2 * e + 1], x1, dq); + } + } + float dot = dq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + dot += __shfl_xor_sync(0xFFFFFFFF, dot, offset); + if (lane == 0) { + dot_slots[par + warp] = dot; + } + __syncthreads(); + ps[2] = 0; + if (ps[0] < ps[1]) { + int first_3 = ps[0]; + int rows_k_1 = ((ps[1] - first_3) < (4) ? (ps[1] - first_3) : (4)); + int sbase_1 = ring + g_step % 3 * 24576; + if (tid == 0) { + mbarrier_arrive_expect_tx(bars_addr + ((int)((unsigned int)g_step % 3)) * 8, rows_k_1 * 6144); + } + if (tir == 0) { + if (rg < rows_k_1) { + int row_g_1 = first_3 + rg; + long long g_srcg_1 = (long long)row_g_1 * g_stride; + long long x_srcg_1 = (long long)row_g_1 * x_stride; + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase_1 + rg * 3072), "l"(reinterpret_cast(g + g_srcg_1)), "r"((unsigned)(3072)), "r"(bars_addr + ((int)((unsigned int)g_step % 3)) * 8), + "l"(0x12F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase_1 + (rg + 4) * 3072), "l"(reinterpret_cast(x + x_srcg_1)), "r"((unsigned)(3072)), "r"(bars_addr + ((int)((unsigned int)g_step % 3)) * 8), + "l"(0x12F0000000000000ULL) : "memory"); + } + } + ps[0] = ps[0] + 4; + ps[2] = (int)(ps[0] >= ps[1]); + if (ps[0] >= ps[1]) { + if (tid == 0) { + int cid_1 = (int)claim[0]; + int nr_3 = rtab[31]; + int r_0_2 = 0; + int walk_3 = 1; + while (walk_3 == 1) { + walk_3 = 0; + if (nr_3 > r_0_2 + 2) { + if (cid_1 >= rtab[(r_0_2 + 1) * 3]) { + r_0_2 += 1; + walk_3 = 1; + } + } + } + if (cid_1 >= rtab[(nr_3 - 1) * 3]) { + rng[0] = T; + rng[1] = T; + } else { + int first_1_1 = rtab[r_0_2 * 3]; + int base_3 = rtab[r_0_2 * 3 + 1]; + int stages_2 = rtab[r_0_2 * 3 + 2]; + rng[0] = (base_3 + (cid_1 - first_1_1) * stages_2) * 4; + int _min_19 = ((rng[0] + stages_2 * 4) < (T) ? (rng[0] + stages_2 * 4) : (T)); + rng[1] = _min_19; + } + int entry_1_2 = (desc_n + 1) % 5; + int word_2_1 = entry_1_2 * 3; + desc[word_2_1] = cid_1; + desc[word_2_1 + 1] = rng[0]; + desc[word_2_1 + 2] = rng[1]; + mbarrier_arrive(dbar_addr + (entry_1_2) * 8); + if (rng[0] < T) { + unsigned int _atomic_fetch_add_20 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_20; + } else { + int nr_0_2 = rtab[31]; + int last_2 = rtab[(nr_0_2 - 1) * 3] + num_bids - 1; + if (T == 0) { + last_2 = 0; + } + if (cid_1 == last_2) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } + } + int entry_0_2 = (desc_n + 1) % 5; + int _parity_dbar_2 = (int)((unsigned int)(desc_n + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_0_2) * 8, _parity_dbar_2); + int word_1_2 = entry_0_2 * 3; + int w0_2_1 = desc[word_1_2]; + int w1_3_1 = desc[word_1_2 + 1]; + int w2_4_1 = desc[word_1_2 + 2]; + int _warp_redux_21; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_21) : "r"(w0_2_1)); + nd[0] = _warp_redux_21; + int _warp_redux_22; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_22) : "r"(w1_3_1)); + nd[1] = _warp_redux_22; + int _warp_redux_23; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_23) : "r"(w2_4_1)); + nd[2] = _warp_redux_23; + ps[0] = nd[1]; + ps[1] = nd[2]; + } + } + int desc_step_1 = ps[2]; + desc_n += desc_step_1; + float tot = 0.0f; + #pragma unroll + for (int kk = 0; kk < 3; kk++) { + float d = dot_slots[par + rg * 3 + kk]; + tot += d; + } + float rq = r_cur; + float cq = __fdiv_rn(tot, 1536.0f); + float kx = rq * rq * rq * cq; + float nkx = -kx; + if (row < ce) { + unsigned long long dx_base = (unsigned long long)row * 1536; + #pragma unroll + for (int v_3 = 0; v_3 < 2; v_3++) { + #pragma unroll + for (int e_1 = 0; e_1 < 4; e_1++) { + unsigned int gwd2 = gw[v_3 * 4 + e_1]; + unsigned int xwd2 = xw[v_3 * 4 + e_1]; + float g2 = __uint_as_float(gwd2 << 16); + float g3 = __uint_as_float(gwd2 & 4294901760u); + float x2 = __uint_as_float(xwd2 << 16); + float x3 = __uint_as_float(xwd2 & 4294901760u); + acc[v_3 * 8 + 2 * e_1] = __fmaf_rn(g2 * x2, rq, acc[v_3 * 8 + 2 * e_1]); + acc[v_3 * 8 + 2 * e_1 + 1] = __fmaf_rn(g3 * x3, rq, acc[v_3 * 8 + 2 * e_1 + 1]); + out[v_3 * 8 + 2 * e_1] = __fmaf_rn(nkx, x2, rq * (g2 * wf[v_3 * 8 + 2 * e_1])); + out[v_3 * 8 + 2 * e_1 + 1] = __fmaf_rn(nkx, x3, rq * (g3 * wf[v_3 * 8 + 2 * e_1 + 1])); + } + int dcol = (v_3 * 96 + tir) * 8; + unsigned int _packed_0[4]; + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 0], (out + v_3 * 8)[_lp*2+1 + 0])); + _packed_0[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 2], (out + v_3 * 8)[_lp*2+1 + 2])); + _packed_0[_lp + 1] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 4], (out + v_3 * 8)[_lp*2+1 + 4])); + _packed_0[_lp + 2] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 6], (out + v_3 * 8)[_lp*2+1 + 6])); + _packed_0[_lp + 3] = *(uint32_t*)&_bf2; + } + reinterpret_cast(dx + (dx_base + (unsigned long long)dcol))[0] = reinterpret_cast(_packed_0)[0]; + } + } + r_cur = r_nxt; + g_step += 1; + if ((k + 1 & 15) == 0) { + #pragma unroll + for (int jkr = 0; jkr < 16; jkr++) { + float ykr = acc[jkr] - cmp[jkr]; + float tkr = hi[jkr] + ykr; + cmp[jkr] = tkr - hi[jkr] - ykr; + hi[jkr] = tkr; + acc[jkr] = 0.0f; + } + } + } + if ((n_steps_c & 15) != 0) { + #pragma unroll + for (int jkt = 0; jkt < 16; jkt++) { + float ykt = acc[jkt] - cmp[jkt]; + float tkt = hi[jkt] + ykt; + cmp[jkt] = tkt - hi[jkt] - ykt; + hi[jkt] = tkt; + acc[jkt] = 0.0f; + } + } + int entry_1_3 = n_cons % 5; + int _parity_dbar_3 = (int)((unsigned int)n_cons / 5 % 2); + mbarrier_wait(dbar_addr + (entry_1_3) * 8, _parity_dbar_3); + int word_1_3 = entry_1_3 * 3; + int w0_1 = desc[word_1_3]; + int w1_1 = desc[word_1_3 + 1]; + int w2_1 = desc[word_1_3 + 2]; + int _warp_redux_24; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_24) : "r"(w0_1)); + nd[0] = _warp_redux_24; + int _warp_redux_25; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_25) : "r"(w1_1)); + nd[1] = _warp_redux_25; + int _warp_redux_26; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_26) : "r"(w2_1)); + nd[2] = _warp_redux_26; + int cid_f = nd[0]; + unsigned long long pbase = (unsigned long long)cid_f * 1536; + #pragma unroll + for (int jk = 0; jk < 16; jk++) { + acc[jk] = hi[jk] - cmp[jk]; + hi[jk] = 0.0f; + cmp[jk] = 0.0f; + } + #pragma unroll + for (int v_4 = 0; v_4 < 2; v_4++) { + int fcol = (v_4 * 96 + tir) * 8; + #pragma unroll + for (int e_2 = 0; e_2 < 8; e_2++) { + fold_sm[rg * 1536 + fcol + e_2] = acc[v_4 * 8 + e_2]; + } + } + __syncthreads(); + float fold_acc[4]; + #pragma unroll + for (int e_3 = 0; e_3 < 4; e_3++) { + fold_acc[e_3] = 0.0f; + } + #pragma unroll + for (int k2 = 0; k2 < 4; k2++) { + #pragma unroll + for (int e_4 = 0; e_4 < 4; e_4++) { + float d2 = fold_sm[k2 * 1536 + tid * 4 + e_4]; + fold_acc[e_4] = fold_acc[e_4] + d2; + } + } + int pcol = tid * 4; + { + float4 _v4 = make_float4(fold_acc[0 + 0], fold_acc[0 + 1], fold_acc[0 + 2], fold_acc[0 + 3]); + *reinterpret_cast(partial + pbase + (unsigned long long)pcol) = _v4; + } + #pragma unroll + for (int j2 = 0; j2 < 16; j2++) { + acc[j2] = 0.0f; + } + int entry_2 = (n_cons + 1) % 5; + int _parity_dbar_4 = (int)((unsigned int)(n_cons + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_2) * 8, _parity_dbar_4); + int word_2_2 = entry_2 * 3; + int w0_2_2 = desc[word_2_2]; + int w1_2 = desc[word_2_2 + 1]; + int w2_2 = desc[word_2_2 + 2]; + int _warp_redux_27; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_27) : "r"(w0_2_2)); + nd[0] = _warp_redux_27; + int _warp_redux_28; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_28) : "r"(w1_2)); + nd[1] = _warp_redux_28; + int _warp_redux_29; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_29) : "r"(w2_2)); + nd[2] = _warp_redux_29; + cb = nd[1]; + ce = nd[2]; + n_cons += 1; + } + + // Cleanup + __syncthreads(); +} + +} // extern "C" diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_binding.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_binding.cu new file mode 100644 index 00000000000..d81db0e3ced --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_binding.cu @@ -0,0 +1,146 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc, for the device source that defines it. Do not edit. +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0); + +namespace cake_binding_kernel_cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckRank(const TensorView& t, const char* name, int rank) { + TVM_FFI_CHECK(t.ndim() >= rank, ValueError) + << KERNEL_NAME << ": '" << name << "' needs at least " << rank << " dims, got " << t.ndim(); +} + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +void Run(TensorView a_g, TensorView a_x, TensorView a_w, TensorView a_r, TensorView a_dx, TensorView a_dw, TensorView a_partial, TensorView a_counters, TensorView a_g_h, int64_t a_T, int64_t a_n_chunks, int64_t a_rows_per_chunk, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_g.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_g, "g"); + check_dtype(a_g, DLDataType{kDLBfloat, 16, 1}, "g"); + check_cuda_tensor(a_x, "x"); + check_dtype(a_x, DLDataType{kDLBfloat, 16, 1}, "x"); + check_cuda_tensor(a_w, "w"); + check_dtype(a_w, DLDataType{kDLBfloat, 16, 1}, "w"); + check_cuda_tensor(a_r, "r"); + check_dtype(a_r, DLDataType{kDLFloat, 32, 1}, "r"); + check_cuda_tensor(a_dx, "dx"); + check_dtype(a_dx, DLDataType{kDLBfloat, 16, 1}, "dx"); + check_cuda_tensor(a_dw, "dw"); + check_dtype(a_dw, DLDataType{kDLFloat, 32, 1}, "dw"); + check_cuda_tensor(a_partial, "partial"); + check_dtype(a_partial, DLDataType{kDLFloat, 32, 1}, "partial"); + check_cuda_tensor(a_counters, "counters"); + check_dtype(a_counters, DLDataType{kDLUInt, 32, 1}, "counters"); + check_cuda_tensor(a_g_h, "g_h"); + check_dtype(a_g_h, DLDataType{kDLBfloat, 16, 1}, "g_h"); + TVM_FFI_CHECK(a_T >= -2147483648LL && a_T <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'T' = " << a_T + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_n_chunks >= -2147483648LL && a_n_chunks <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'n_chunks' = " << a_n_chunks + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_rows_per_chunk >= -2147483648LL && a_rows_per_chunk <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'rows_per_chunk' = " << a_rows_per_chunk + << " is outside the i32 range [-2147483648, 2147483647]"; + CheckRank(a_g, "g", 1); + CheckUnitInnerStride(a_g, "g"); + CheckRank(a_x, "x", 1); + CheckUnitInnerStride(a_x, "x"); + check_same_device(a_x, a_g, "x", "g"); + check_contiguous(a_w, "w"); + check_same_device(a_w, a_g, "w", "g"); + check_contiguous(a_r, "r"); + check_same_device(a_r, a_g, "r", "g"); + check_contiguous(a_dx, "dx"); + check_same_device(a_dx, a_g, "dx", "g"); + check_contiguous(a_dw, "dw"); + check_same_device(a_dw, a_g, "dw", "g"); + check_contiguous(a_partial, "partial"); + check_same_device(a_partial, a_g, "partial", "g"); + check_contiguous(a_counters, "counters"); + check_same_device(a_counters, a_g, "counters", "g"); + CheckRank(a_g_h, "g_h", 1); + CheckUnitInnerStride(a_g_h, "g_h"); + check_same_device(a_g_h, a_g, "g_h", "g"); + TVM_FFI_CHECK((a_n_chunks >= 1LL), ValueError) + << KERNEL_NAME << ": requires (n_chunks >= 1)"; + TVM_FFI_CHECK((a_rows_per_chunk >= 1LL), ValueError) + << KERNEL_NAME << ": requires (rows_per_chunk >= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc), cudaFuncAttributeMaxDynamicSharedMemorySize, 99840), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + void* v_g = DataPtr(a_g); + void* v_x = DataPtr(a_x); + void* v_w = DataPtr(a_w); + void* v_r = DataPtr(a_r); + void* v_dx = DataPtr(a_dx); + void* v_dw = DataPtr(a_dw); + void* v_partial = DataPtr(a_partial); + void* v_counters = DataPtr(a_counters); + void* v_g_h = DataPtr(a_g_h); + int32_t v_T = static_cast(a_T); + int32_t v_n_chunks = static_cast(a_n_chunks); + int32_t v_rows_per_chunk = static_cast(a_rows_per_chunk); + int64_t v_g_stride_p0 = a_g.stride(0); + int64_t v_x_stride_p0 = a_x.stride(0); + int64_t v_g_h_stride_p0 = a_g_h.stride(0); + void* kernel_args[] = {&v_g, &v_x, &v_w, &v_r, &v_dx, &v_dw, &v_partial, &v_counters, &v_g_h, &v_T, &v_n_chunks, &v_rows_per_chunk, &v_g_stride_p0, &v_x_stride_p0, &v_g_h_stride_p0}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(384, 1, 1); + config.dynamicSmemBytes = 99840; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_kernel.cu b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_kernel.cu new file mode 100644 index 00000000000..8be580f881d --- /dev/null +++ b/flashinfer/cake_rmsnorm_det_train_backend/csrc/cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc_kernel.cu @@ -0,0 +1,826 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define SMEM_STAGE_BUF_OFF 1024 +#define SMEM_STAGE_BUF_STAGE_BYTES 73728 +#define SMEM_STAGE_BUF_STRIDE 73728 +#define SMEM_ROW_FOLD_OFF 1024 +#define SMEM_ROW_FOLD_STAGE_BYTES 24576 +#define SMEM_ROW_FOLD_STRIDE 24576 +#define SMEM_RED_FOLD_OFF 25600 +#define SMEM_RED_FOLD_STAGE_BYTES 6144 +#define SMEM_RED_FOLD_STRIDE 6144 +#define SMEM_DOT_SLOTS_OFF 74752 +#define SMEM_DOT_SLOTS_STAGE_BYTES 96 +#define SMEM_DOT_SLOTS_STRIDE 96 +#define SMEM_PART_SM_OFF 74880 +#define SMEM_PART_SM_STAGE_BYTES 16 +#define SMEM_PART_SM_STRIDE 16 +#define SMEM_DESC_OFF 99584 +#define SMEM_DESC_STAGE_BYTES 60 +#define SMEM_DESC_STRIDE 60 +#define SMEM_RTAB_OFF 99712 +#define SMEM_RTAB_STAGE_BYTES 128 +#define SMEM_RTAB_STRIDE 128 +#define SMEM_FOLD_SM_OFF 75008 +#define SMEM_FOLD_SM_STAGE_BYTES 24576 +#define SMEM_FOLD_SM_STRIDE 24576 +#define SMEM_TOTAL 99840 +#define THREADS 384 +#define LAUNCH_MIN_BLOCKS 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_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"); +} + + + +extern "C" { + +__global__ __launch_bounds__(384, LAUNCH_MIN_BLOCKS) void +kernel_cake_rmsnorm_det_train_e694b0a25ea01d1f6cdc(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define bars_addr (mbar_base + 0) + #define dbar_addr (mbar_base + 96) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + uint8_t* stage_buf = reinterpret_cast(smem_raw + 1024); + const int stage_buf_addr = smem + 1024; + float* row_fold = reinterpret_cast(smem_raw + 1024); + const int row_fold_addr = smem + 1024; + float* red_fold = reinterpret_cast(smem_raw + 25600); + const int red_fold_addr = smem + 25600; + float* dot_slots = reinterpret_cast(smem_raw + 74752); + const int dot_slots_addr = smem + 74752; + float* part_sm = reinterpret_cast(smem_raw + 74880); + const int part_sm_addr = smem + 74880; + int* desc = reinterpret_cast(smem_raw + 99584); + const int desc_addr = smem + 99584; + int* rtab = reinterpret_cast(smem_raw + 99712); + const int rtab_addr = smem + 99712; + float* fold_sm = reinterpret_cast(smem_raw + 75008); + const int fold_sm_addr = smem + 75008; + + // Mbarrier init (2 pipeline groups, 0 ordered-sequence groups, 17 barriers) + // Mbarriers at smem_raw[0..136) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // bars: 12 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + mbarrier_init(smem + 8, 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); + // dbar: 5 barriers, init_count=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); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncthreads(); + + // === Task calls (declaration order) === + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + long long g_stride = g_stride_p0; + long long x_stride = x_stride_p0; + long long gh_stride = g_h_stride_p0; + int c = bid; + int rg = tid / 96; + int tir = tid - rg * 96; + unsigned int claim[1]; + if (tid == 0) { + unsigned int _atomic_fetch_add_0 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_0; + int S = (T + 3) / 4; + int _min_1 = ((rows_per_chunk) < (S) ? (rows_per_chunk) : (S)); + int sl = ((_min_1) > (1) ? (_min_1) : (1)); + int share = (S + sl - 1) / sl; + int tail = (((share + 11) / 12) > (1) ? ((share + 11) / 12) : (1)); + int size = (share + 1) / 2; + int rem = S; + int base = 0; + int first = 0; + int nr = 0; + int walk = 1; + while (walk == 1) { + walk = 0; + if (nr < 5) { + if (size > tail) { + if (rem >= sl * size) { + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = size; + first += sl; + base += sl * size; + rem -= sl * size; + size = (size + 1) / 2; + nr += 1; + walk = 1; + } + } + } + } + int q = rem / sl; + int m = rem - q * sl; + if (m > 0) { + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = q + 1; + first += m; + base += m * (q + 1); + nr += 1; + } + if (q > 0) { + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = q; + first += sl - m; + base += (sl - m) * q; + nr += 1; + } + rtab[nr * 3] = first; + rtab[nr * 3 + 1] = base; + rtab[nr * 3 + 2] = 0; + rtab[31] = nr + 1; + } + int ring = (int)stage_buf_addr; + int my_g = ring + rg * 3072 + tir * 16; + unsigned int ww[8]; + #pragma unroll + for (int v = 0; v < 2; v++) { + int wcol = (v * 96 + tir) * 8; + { + const uint4* _vptr_2 = reinterpret_cast(w + (unsigned long long)wcol); + uint4* _vdst_2 = reinterpret_cast(&ww[v * 4]); + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vdst_2[_blk] = _vptr_2[_blk]; + } + } + } + int rng[2]; + int nd[3]; + if (tid == 0) { + int cid0 = (int)claim[0]; + int nr_1 = rtab[31]; + int r_0 = 0; + int walk_1 = 1; + while (walk_1 == 1) { + walk_1 = 0; + if (nr_1 > r_0 + 2) { + if (cid0 >= rtab[(r_0 + 1) * 3]) { + r_0 += 1; + walk_1 = 1; + } + } + } + if (cid0 >= rtab[(nr_1 - 1) * 3]) { + rng[0] = T; + rng[1] = T; + } else { + int first_1 = rtab[r_0 * 3]; + int base_1 = rtab[r_0 * 3 + 1]; + int stages = rtab[r_0 * 3 + 2]; + rng[0] = (base_1 + (cid0 - first_1) * stages) * 4; + int _min_3 = ((rng[0] + stages * 4) < (T) ? (rng[0] + stages * 4) : (T)); + rng[1] = _min_3; + } + int entry = 0; + int word = entry * 3; + desc[word] = cid0; + desc[word + 1] = rng[0]; + desc[word + 2] = rng[1]; + mbarrier_arrive(dbar_addr + (entry) * 8); + if (rng[0] < T) { + unsigned int _atomic_fetch_add_4 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_4; + } else { + int nr_0 = rtab[31]; + int last = rtab[(nr_0 - 1) * 3] + num_bids - 1; + if (T == 0) { + last = 0; + } + if (cid0 == last) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } + } + int entry_1 = 0; + mbarrier_wait(dbar_addr + (entry_1) * 8, 0); + int word_1 = entry_1 * 3; + int w0 = desc[word_1]; + int w1 = desc[word_1 + 1]; + int w2 = desc[word_1 + 2]; + int _warp_redux_5; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_5) : "r"(w0)); + nd[0] = _warp_redux_5; + int _warp_redux_6; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_6) : "r"(w1)); + nd[1] = _warp_redux_6; + int _warp_redux_7; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_7) : "r"(w2)); + nd[2] = _warp_redux_7; + int c0 = nd[0]; + int cb = nd[1]; + int ce = nd[2]; + int ps[3]; + ps[0] = nd[1]; + ps[1] = nd[2]; + int desc_n = 0; + #pragma unroll + for (int s = 0; s < 3; s++) { + ps[2] = 0; + if (ps[0] < ps[1]) { + int first_2 = ps[0]; + int rows_k = ((ps[1] - first_2) < (4) ? (ps[1] - first_2) : (4)); + int sbase = ring + s * 24576; + int gstage = rg + 4 * s; + if (tir == 0) { + if (rg < rows_k) { + mbarrier_arrive_expect_tx(bars_addr + (gstage) * 8, 6144); + int row_g = first_2 + rg; + long long g_srcg = (long long)row_g * g_stride; + long long x_srcg = (long long)row_g * x_stride; + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase + rg * 3072), "l"(reinterpret_cast(g + g_srcg)), "r"((unsigned)(3072)), "r"(bars_addr + (gstage) * 8), + "l"(0x14F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase + (rg + 4) * 3072), "l"(reinterpret_cast(x + x_srcg)), "r"((unsigned)(3072)), "r"(bars_addr + (gstage) * 8), + "l"(0x14F0000000000000ULL) : "memory"); + } else { + mbarrier_arrive(bars_addr + (gstage) * 8); + } + } + ps[0] = ps[0] + 4; + ps[2] = (int)(ps[0] >= ps[1]); + if (ps[0] >= ps[1]) { + if (tid == 0) { + int cid = (int)claim[0]; + int nr_2 = rtab[31]; + int r_0_1 = 0; + int walk_2 = 1; + while (walk_2 == 1) { + walk_2 = 0; + if (nr_2 > r_0_1 + 2) { + if (cid >= rtab[(r_0_1 + 1) * 3]) { + r_0_1 += 1; + walk_2 = 1; + } + } + } + if (cid >= rtab[(nr_2 - 1) * 3]) { + rng[0] = T; + rng[1] = T; + } else { + int first_0 = rtab[r_0_1 * 3]; + int base_2 = rtab[r_0_1 * 3 + 1]; + int stages_1 = rtab[r_0_1 * 3 + 2]; + rng[0] = (base_2 + (cid - first_0) * stages_1) * 4; + int _min_8 = ((rng[0] + stages_1 * 4) < (T) ? (rng[0] + stages_1 * 4) : (T)); + rng[1] = _min_8; + } + int entry_1_1 = (desc_n + 1) % 5; + int word_2 = entry_1_1 * 3; + desc[word_2] = cid; + desc[word_2 + 1] = rng[0]; + desc[word_2 + 2] = rng[1]; + mbarrier_arrive(dbar_addr + (entry_1_1) * 8); + if (rng[0] < T) { + unsigned int _atomic_fetch_add_9 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_9; + } else { + int nr_0_1 = rtab[31]; + int last_1 = rtab[(nr_0_1 - 1) * 3] + num_bids - 1; + if (T == 0) { + last_1 = 0; + } + if (cid == last_1) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } + } + int entry_0 = (desc_n + 1) % 5; + int _parity_dbar = (int)((unsigned int)(desc_n + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_0) * 8, _parity_dbar); + int word_1_1 = entry_0 * 3; + int w0_2 = desc[word_1_1]; + int w1_3 = desc[word_1_1 + 1]; + int w2_4 = desc[word_1_1 + 2]; + int _warp_redux_10; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_10) : "r"(w0_2)); + nd[0] = _warp_redux_10; + int _warp_redux_11; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_11) : "r"(w1_3)); + nd[1] = _warp_redux_11; + int _warp_redux_12; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_12) : "r"(w2_4)); + nd[2] = _warp_redux_12; + ps[0] = nd[1]; + ps[1] = nd[2]; + } + } + int desc_step = ps[2]; + desc_n += desc_step; + } + float wf[16]; + #pragma unroll + for (int p = 0; p < 8; p++) { + unsigned int wwd = ww[p]; + wf[2 * p] = __uint_as_float(wwd << 16); + wf[2 * p + 1] = __uint_as_float(wwd & 4294901760u); + } + float acc[16]; + #pragma unroll + for (int j = 0; j < 16; j++) { + acc[j] = 0.0f; + } + float cmp[16]; + #pragma unroll + for (int jc = 0; jc < 16; jc++) { + cmp[jc] = 0.0f; + } + float hi[16]; + #pragma unroll + for (int jh = 0; jh < 16; jh++) { + hi[jh] = 0.0f; + } + int n_cons = 0; + int g_step = 0; + float r_cur = 0.0f; + if (cb < ce) { + int row0 = ((cb + rg) < (ce - 1) ? (cb + rg) : (ce - 1)); + r_cur = r[row0]; + } + unsigned int gw[8]; + unsigned int xw[8]; + float out[16]; + if (cb >= ce) { + if (c0 < n_chunks) { + unsigned long long pbase0 = (unsigned long long)c0 * 1536; + float fold0[4]; + #pragma unroll + for (int e0 = 0; e0 < 4; e0++) { + fold0[e0] = 0.0f; + } + int pcol0 = tid * 4; + { + float4 _v4 = make_float4(fold0[0 + 0], fold0[0 + 1], fold0[0 + 2], fold0[0 + 3]); + *reinterpret_cast(partial + pbase0 + (unsigned long long)pcol0) = _v4; + } + } + } + while (cb < ce) { + int n_steps_c = (ce - cb + 3) / 4; + #pragma unroll 1 + for (int k = 0; k < n_steps_c; k++) { + int slot = g_step % 3; + int par = (g_step & 1) * 12; + int row = cb + k * 4 + rg; + int row_n = ((row + 4) < (ce - 1) ? (row + 4) : (ce - 1)); + if (n_steps_c <= k + 1) { + int entry_0_1 = (n_cons + 1) % 5; + int _parity_dbar_1 = (int)((unsigned int)(n_cons + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_0_1) * 8, _parity_dbar_1); + int word_0 = entry_0_1 * 3; + int w0_0 = desc[word_0]; + int w1_0 = desc[word_0 + 1]; + int w2_0 = desc[word_0 + 2]; + int _warp_redux_13; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_13) : "r"(w0_0)); + nd[0] = _warp_redux_13; + int _warp_redux_14; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_14) : "r"(w1_0)); + nd[1] = _warp_redux_14; + int _warp_redux_15; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_15) : "r"(w2_0)); + nd[2] = _warp_redux_15; + row_n = cb; + if (nd[1] < T) { + int _min_16 = ((nd[1] + rg) < (nd[2] - 1) ? (nd[1] + rg) : (nd[2] - 1)); + row_n = _min_16; + } + } + float r_nxt = r[row_n]; + int _stage_bars = rg + 4 * (int)((unsigned int)g_step % 3); + int _parity_bars = (int)((unsigned int)g_step / 3 % 2); + mbarrier_wait(bars_addr + (_stage_bars) * 8, _parity_bars); + int sg = my_g + slot * 24576; + #pragma unroll + for (int v_1 = 0; v_1 < 2; v_1++) { + __int128_t _smem_raw_17; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_17) : "r"(sg + v_1 * 1536)); + gw[v_1 * 4] = reinterpret_cast(&_smem_raw_17)[0]; + gw[(v_1 * 4) + 1] = reinterpret_cast(&_smem_raw_17)[1]; + gw[(v_1 * 4) + 2] = reinterpret_cast(&_smem_raw_17)[2]; + gw[(v_1 * 4) + 3] = reinterpret_cast(&_smem_raw_17)[3]; + __int128_t _smem_raw_18; + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(_smem_raw_18) : "r"(sg + (12288 + v_1 * 1536))); + xw[v_1 * 4] = reinterpret_cast(&_smem_raw_18)[0]; + xw[(v_1 * 4) + 1] = reinterpret_cast(&_smem_raw_18)[1]; + xw[(v_1 * 4) + 2] = reinterpret_cast(&_smem_raw_18)[2]; + xw[(v_1 * 4) + 3] = reinterpret_cast(&_smem_raw_18)[3]; + } + float dq = 0.0f; + #pragma unroll + for (int v_2 = 0; v_2 < 2; v_2++) { + #pragma unroll + for (int e = 0; e < 4; e++) { + unsigned int gwd = gw[v_2 * 4 + e]; + unsigned int xwd = xw[v_2 * 4 + e]; + float g0 = __uint_as_float(gwd << 16); + float g1 = __uint_as_float(gwd & 4294901760u); + float x0 = __uint_as_float(xwd << 16); + float x1 = __uint_as_float(xwd & 4294901760u); + dq = __fmaf_rn(g0 * wf[v_2 * 8 + 2 * e], x0, dq); + dq = __fmaf_rn(g1 * wf[v_2 * 8 + 2 * e + 1], x1, dq); + } + } + float dot = dq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + dot += __shfl_xor_sync(0xFFFFFFFF, dot, offset); + if (lane == 0) { + dot_slots[par + warp] = dot; + } + asm volatile("barrier.sync %0, 96;" :: "r"(1 + rg) : "memory"); + ps[2] = 0; + if (ps[0] < ps[1]) { + int first_3 = ps[0]; + int rows_k_1 = ((ps[1] - first_3) < (4) ? (ps[1] - first_3) : (4)); + int sbase_1 = ring + g_step % 3 * 24576; + int gstage_1 = rg + 4 * (g_step % 3); + if (tir == 0) { + if (rg < rows_k_1) { + mbarrier_arrive_expect_tx(bars_addr + (gstage_1) * 8, 6144); + int row_g_1 = first_3 + rg; + long long g_srcg_1 = (long long)row_g_1 * g_stride; + long long x_srcg_1 = (long long)row_g_1 * x_stride; + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase_1 + rg * 3072), "l"(reinterpret_cast(g + g_srcg_1)), "r"((unsigned)(3072)), "r"(bars_addr + (gstage_1) * 8), + "l"(0x14F0000000000000ULL) : "memory"); + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1], %2, [%3], %4;" + :: "r"(sbase_1 + (rg + 4) * 3072), "l"(reinterpret_cast(x + x_srcg_1)), "r"((unsigned)(3072)), "r"(bars_addr + (gstage_1) * 8), + "l"(0x14F0000000000000ULL) : "memory"); + } else { + mbarrier_arrive(bars_addr + (gstage_1) * 8); + } + } + ps[0] = ps[0] + 4; + ps[2] = (int)(ps[0] >= ps[1]); + if (ps[0] >= ps[1]) { + if (tid == 0) { + int cid_1 = (int)claim[0]; + int nr_3 = rtab[31]; + int r_0_2 = 0; + int walk_3 = 1; + while (walk_3 == 1) { + walk_3 = 0; + if (nr_3 > r_0_2 + 2) { + if (cid_1 >= rtab[(r_0_2 + 1) * 3]) { + r_0_2 += 1; + walk_3 = 1; + } + } + } + if (cid_1 >= rtab[(nr_3 - 1) * 3]) { + rng[0] = T; + rng[1] = T; + } else { + int first_1_1 = rtab[r_0_2 * 3]; + int base_3 = rtab[r_0_2 * 3 + 1]; + int stages_2 = rtab[r_0_2 * 3 + 2]; + rng[0] = (base_3 + (cid_1 - first_1_1) * stages_2) * 4; + int _min_19 = ((rng[0] + stages_2 * 4) < (T) ? (rng[0] + stages_2 * 4) : (T)); + rng[1] = _min_19; + } + int entry_1_2 = (desc_n + 1) % 5; + int word_2_1 = entry_1_2 * 3; + desc[word_2_1] = cid_1; + desc[word_2_1 + 1] = rng[0]; + desc[word_2_1 + 2] = rng[1]; + mbarrier_arrive(dbar_addr + (entry_1_2) * 8); + if (rng[0] < T) { + unsigned int _atomic_fetch_add_20 = atomicAdd(&counters[0], 1); + claim[0] = _atomic_fetch_add_20; + } else { + int nr_0_2 = rtab[31]; + int last_2 = rtab[(nr_0_2 - 1) * 3] + num_bids - 1; + if (T == 0) { + last_2 = 0; + } + if (cid_1 == last_2) { + *(reinterpret_cast(counters) + (0)) = 0; + } + } + } + int entry_0_2 = (desc_n + 1) % 5; + int _parity_dbar_2 = (int)((unsigned int)(desc_n + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_0_2) * 8, _parity_dbar_2); + int word_1_2 = entry_0_2 * 3; + int w0_2_1 = desc[word_1_2]; + int w1_3_1 = desc[word_1_2 + 1]; + int w2_4_1 = desc[word_1_2 + 2]; + int _warp_redux_21; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_21) : "r"(w0_2_1)); + nd[0] = _warp_redux_21; + int _warp_redux_22; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_22) : "r"(w1_3_1)); + nd[1] = _warp_redux_22; + int _warp_redux_23; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_23) : "r"(w2_4_1)); + nd[2] = _warp_redux_23; + ps[0] = nd[1]; + ps[1] = nd[2]; + } + } + int desc_step_1 = ps[2]; + desc_n += desc_step_1; + float tot = 0.0f; + #pragma unroll + for (int kk = 0; kk < 3; kk++) { + float d = dot_slots[par + rg * 3 + kk]; + tot += d; + } + float rq = r_cur; + float cq = __fdiv_rn(tot, 1536.0f); + float kx = rq * rq * rq * cq; + float nkx = -kx; + if (row < ce) { + unsigned long long dx_base = (unsigned long long)row * 1536; + #pragma unroll + for (int v_3 = 0; v_3 < 2; v_3++) { + #pragma unroll + for (int e_1 = 0; e_1 < 4; e_1++) { + unsigned int gwd2 = gw[v_3 * 4 + e_1]; + unsigned int xwd2 = xw[v_3 * 4 + e_1]; + float g2 = __uint_as_float(gwd2 << 16); + float g3 = __uint_as_float(gwd2 & 4294901760u); + float x2 = __uint_as_float(xwd2 << 16); + float x3 = __uint_as_float(xwd2 & 4294901760u); + acc[v_3 * 8 + 2 * e_1] = __fmaf_rn(g2 * x2, rq, acc[v_3 * 8 + 2 * e_1]); + acc[v_3 * 8 + 2 * e_1 + 1] = __fmaf_rn(g3 * x3, rq, acc[v_3 * 8 + 2 * e_1 + 1]); + out[v_3 * 8 + 2 * e_1] = __fmaf_rn(nkx, x2, rq * (g2 * wf[v_3 * 8 + 2 * e_1])); + out[v_3 * 8 + 2 * e_1 + 1] = __fmaf_rn(nkx, x3, rq * (g3 * wf[v_3 * 8 + 2 * e_1 + 1])); + } + int dcol = (v_3 * 96 + tir) * 8; + unsigned int _packed_0[4]; + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 0], (out + v_3 * 8)[_lp*2+1 + 0])); + _packed_0[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 2], (out + v_3 * 8)[_lp*2+1 + 2])); + _packed_0[_lp + 1] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 4], (out + v_3 * 8)[_lp*2+1 + 4])); + _packed_0[_lp + 2] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int _lp = 0; _lp < 1; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2((out + v_3 * 8)[_lp*2 + 6], (out + v_3 * 8)[_lp*2+1 + 6])); + _packed_0[_lp + 3] = *(uint32_t*)&_bf2; + } + reinterpret_cast(dx + (dx_base + (unsigned long long)dcol))[0] = reinterpret_cast(_packed_0)[0]; + } + } + r_cur = r_nxt; + g_step += 1; + if ((k + 1 & 3) == 0) { + #pragma unroll + for (int jkr = 0; jkr < 16; jkr++) { + float ykr = acc[jkr] - cmp[jkr]; + float tkr = hi[jkr] + ykr; + cmp[jkr] = tkr - hi[jkr] - ykr; + hi[jkr] = tkr; + acc[jkr] = 0.0f; + } + } + } + if ((n_steps_c & 3) != 0) { + #pragma unroll + for (int jkt = 0; jkt < 16; jkt++) { + float ykt = acc[jkt] - cmp[jkt]; + float tkt = hi[jkt] + ykt; + cmp[jkt] = tkt - hi[jkt] - ykt; + hi[jkt] = tkt; + acc[jkt] = 0.0f; + } + } + int entry_1_3 = n_cons % 5; + int _parity_dbar_3 = (int)((unsigned int)n_cons / 5 % 2); + mbarrier_wait(dbar_addr + (entry_1_3) * 8, _parity_dbar_3); + int word_1_3 = entry_1_3 * 3; + int w0_1 = desc[word_1_3]; + int w1_1 = desc[word_1_3 + 1]; + int w2_1 = desc[word_1_3 + 2]; + int _warp_redux_24; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_24) : "r"(w0_1)); + nd[0] = _warp_redux_24; + int _warp_redux_25; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_25) : "r"(w1_1)); + nd[1] = _warp_redux_25; + int _warp_redux_26; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_26) : "r"(w2_1)); + nd[2] = _warp_redux_26; + int cid_f = nd[0]; + unsigned long long pbase = (unsigned long long)cid_f * 1536; + __syncthreads(); + #pragma unroll + for (int jk = 0; jk < 16; jk++) { + acc[jk] = hi[jk] - cmp[jk]; + hi[jk] = 0.0f; + cmp[jk] = 0.0f; + } + #pragma unroll + for (int v_4 = 0; v_4 < 2; v_4++) { + int fcol = (v_4 * 96 + tir) * 8; + #pragma unroll + for (int e_2 = 0; e_2 < 8; e_2++) { + fold_sm[rg * 1536 + fcol + e_2] = acc[v_4 * 8 + e_2]; + } + } + __syncthreads(); + float fold_acc[4]; + #pragma unroll + for (int e_3 = 0; e_3 < 4; e_3++) { + fold_acc[e_3] = 0.0f; + } + #pragma unroll + for (int k2 = 0; k2 < 4; k2++) { + #pragma unroll + for (int e_4 = 0; e_4 < 4; e_4++) { + float d2 = fold_sm[k2 * 1536 + tid * 4 + e_4]; + fold_acc[e_4] = fold_acc[e_4] + d2; + } + } + int pcol = tid * 4; + { + float4 _v4 = make_float4(fold_acc[0 + 0], fold_acc[0 + 1], fold_acc[0 + 2], fold_acc[0 + 3]); + *reinterpret_cast(partial + pbase + (unsigned long long)pcol) = _v4; + } + #pragma unroll + for (int j2 = 0; j2 < 16; j2++) { + acc[j2] = 0.0f; + } + int entry_2 = (n_cons + 1) % 5; + int _parity_dbar_4 = (int)((unsigned int)(n_cons + 1) / 5 % 2); + mbarrier_wait(dbar_addr + (entry_2) * 8, _parity_dbar_4); + int word_2_2 = entry_2 * 3; + int w0_2_2 = desc[word_2_2]; + int w1_2 = desc[word_2_2 + 1]; + int w2_2 = desc[word_2_2 + 2]; + int _warp_redux_27; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_27) : "r"(w0_2_2)); + nd[0] = _warp_redux_27; + int _warp_redux_28; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_28) : "r"(w1_2)); + nd[1] = _warp_redux_28; + int _warp_redux_29; + asm volatile("redux.sync.max.s32 %0, %1, 0xffffffff;" : "=r"(_warp_redux_29) : "r"(w2_2)); + nd[2] = _warp_redux_29; + cb = nd[1]; + ce = nd[2]; + n_cons += 1; + } + + // Cleanup + __syncthreads(); +} + +} // extern "C" diff --git a/tests/norm/test_cake_rmsnorm_det_train_pairs.py b/tests/norm/test_cake_rmsnorm_det_train_pairs.py index 056e63d184c..e40a9910426 100644 --- a/tests/norm/test_cake_rmsnorm_det_train_pairs.py +++ b/tests/norm/test_cake_rmsnorm_det_train_pairs.py @@ -40,7 +40,9 @@ SUPPORTED_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) HIDDEN_SIZES = (6144, 5120, 4096, 2048, 1536, 512, 256) REDUCE_ROLE = "bwd_reduce" -SHIPPED_SLABS = (16, 32, 64) +# The column widths of the dw reduce slabs, one CTA each: 4 for the column-split H = 256 reduce (one column of a +# four-column slab per lane). +SHIPPED_SLABS = (4, 16, 32, 64) # Row counts of the contract: the small counts, both sides of every backward window edge of every route # (``test_rows_cover_every_backward_window_edge``) and the training shapes. ROWS = ( @@ -317,12 +319,14 @@ def test_rows_cover_every_backward_window_edge() -> None: # five-CTA ring at issue distance three without the L2 hint (``a3``), both at ``tail_div=4``, in two sessions per # architecture (BF16 and FP32 weight alike); sm_107a runs the pipeline to 65536 rows and the ring above. Above the # ring's dw accuracy edge of each architecture (80896 rows on sm_100a, 99840 on sm_103a, 70656 on sm_107a) the ring -# with compensated dw sums (``dc``) runs instead, and above 1572864 rows its form with hi / lo partial rows +# with compensated dw sums (``dc``; on sm_100a and sm_107a its lean form ``dcln``, two rows per iteration with the +# same arithmetic in the same order) runs instead, and above 1572864 rows its form with hi / lo partial rows # (``dchl``). H256_BWD_PROGRAMS = { "pd4": "rmsnorm_train_bwd_h256_t128x4_pd4_k_s32_split_pts_cl_td4", "a3": "rmsnorm_train_bwd_h256_t128x5_cpa8a3nhkc8_k_s32_split_pts_cl_td4_cs", "dc": "rmsnorm_train_bwd_h256_t128x5_cpa8a3nhftdf_k_s32_split_pts_cl_td4_cs", + "dcln": "rmsnorm_train_bwd_h256_t128x5_cpa8a3nhftdfln_k_s32_split_pts_cl_td4_cs", "dchl": "rmsnorm_train_bwd_h256_t128x5_cpa8a3nhftdfhl_k_s32_split_pts_cl_td4_cs", } H256_BWD_SPLIT = { @@ -333,18 +337,18 @@ def test_rows_cover_every_backward_window_edge() -> None: (5120, 12288, "a3"), (12288, 40960, "pd4"), (40960, edge, "a3"), - (edge, hl, "dc"), + (edge, hl, dc), (hl, None, "dchl"), ) - for arch, edge, hl in ( - ("sm_100a", 80896, 1572864), - ("sm_103a", 99840, 1572864), + for arch, edge, hl, dc in ( + ("sm_100a", 80896, 1572864, "dcln"), + ("sm_103a", 99840, 1572864, "dc"), ) }, "sm_107a": ( (0, 65536, "pd4"), (65536, 70656, "a3"), - (70656, 1572864, "dc"), + (70656, 1572864, "dcln"), (1572864, None, "dchl"), ), } From d409c377930e7dfe1d0dfb79ecb3b904ea1627d8 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Mon, 2 Nov 2026 05:53:44 -0800 Subject: [PATCH 2/2] perf(cake_rmsnorm_det): keep the shared-memory compensated H 1536 backward for the BF16 weight on SM103 On SM103 above 40448 rows (H = 1536) the register-compensated backward is faster than the shared-memory compensated one with the FP32 weight on every GPU measured, and not with the BF16 weight. The BF16-weight route returns to the shared-memory program; the FP32-weight route keeps the FP32-weight form of the register program (the same window, launch, workspace and dw reduce, bitwise equal dx and dw). cake_windows.py lists that FP32-weight route with its own windows and edges and records the two SM107 H 256 backward edges measured since (70656 and 1572864 rows: y, rstd and dx bitwise equal across them); the pairs test knows such routes. The kernels are unchanged. --- .../cake_jit.py | 2 +- .../cake_windows.py | 23 ++++++++----------- .../norm/test_cake_rmsnorm_det_train_pairs.py | 16 +++++++++---- 3 files changed, 22 insertions(+), 19 deletions(-) diff --git a/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py b/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py index c06fea652d5..f5fb3b0d85f 100644 --- a/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py +++ b/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py @@ -12116,7 +12116,7 @@ "cake_rmsnorm_det_train_1df7db17277d363b5168", "cake_rmsnorm_det_train_4df08ce51fb01ba1193f", "cake_rmsnorm_det_train_f522f5b8d4343d96ca04", - "cake_rmsnorm_det_train_16a5e0b1afd90f38ed44", + "cake_rmsnorm_det_train_cbeae168d70edd196417", "cake_rmsnorm_det_train_9b727c2c9933b2ca6724", "cake_rmsnorm_det_train_1863320f3158d35eead6", "cake_rmsnorm_det_train_9b727c2c9933b2ca6724__4096_14928", diff --git a/flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py b/flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py index d5b9dca63f3..05267bb43ab 100644 --- a/flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py +++ b/flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py @@ -237,7 +237,7 @@ (4096, 14928, "cake_rmsnorm_det_train_1df7db17277d363b5168"), (14928, 15566, "cake_rmsnorm_det_train_4df08ce51fb01ba1193f"), (15566, 40448, "cake_rmsnorm_det_train_f522f5b8d4343d96ca04"), - (40448, None, "cake_rmsnorm_det_train_16a5e0b1afd90f38ed44"), + (40448, None, "cake_rmsnorm_det_train_cbeae168d70edd196417"), ), "bwd_w32_h1536": ( (0, 1024, "cake_rmsnorm_det_train_baff9ecbafe7de1e6030"), @@ -808,16 +808,8 @@ (16384, "same"), (50688, "same"), ), - "bwd_h256": ( - (65536, "unmeasured"), - (70656, "unmeasured"), - (1572864, "unmeasured"), - ), - "bwd_w32_h256": ( - (65536, "unmeasured"), - (70656, "unmeasured"), - (1572864, "unmeasured"), - ), + "bwd_h256": ((65536, "unmeasured"), (70656, "same"), (1572864, "same")), + "bwd_w32_h256": ((65536, "unmeasured"), (70656, "same"), (1572864, "same")), "bwd_h4096": ((256, "same"), (1024, "same"), (4096, "same"), (24064, "same")), "bwd_w32_h4096": ( (256, "same"), @@ -1157,7 +1149,10 @@ * ``sm_100a`` ``bwd_residual_h6144``: (0, inf) 0dcdc563 * ``sm_103a`` ``bwd_h1536``: (0, 1024] 7dc29629 = (1024, 4096] c1a9efdf = (4096, 14928] 1df7db17 = (14928, 15566] 4df08ce5 = (15566, 40448] f522f5b8 = - (40448, inf) 16a5e0b1 + (40448, inf) cbeae168 +* ``sm_103a`` ``bwd_w32_h1536``: (0, 1024] baff9ecb = (1024, 4096] 17d9cdfd = + (4096, 14928] a99e2b38 = (14928, 15566] 7a72e97b = (15566, 40448] 109f19bb = + (40448, inf) ca6211f3 * ``sm_103a`` ``bwd_h2048``: (0, 3072] caf6978e = (3072, 4096] 9ac8d21e ~ (4096, 16384] caf6978e = (16384, 32768] caf6978e = (32768, 55296] 302f64db = (55296, inf) 018f0cd7 @@ -1176,8 +1171,8 @@ (16384, 45056] 56d26996 = (45056, inf) e694b0a2 * ``sm_107a`` ``bwd_h2048``: (0, 256] 84efa2af = (256, 1024] 845cdcd8 ~ (1024, 16384] 81c3d8ab = (16384, 50688] 845cdcd8 = (50688, inf) 694decfb -* ``sm_107a`` ``bwd_h256``: (0, 65536] af005462 ? (65536, 70656] a385ae16 ? - (70656, 1572864] 083ea823 ? (1572864, inf) c18297d7 +* ``sm_107a`` ``bwd_h256``: (0, 65536] af005462 ? (65536, 70656] a385ae16 = + (70656, 1572864] 083ea823 = (1572864, inf) c18297d7 * ``sm_107a`` ``bwd_h4096``: (0, 256] 0da5dd4b = (256, 1024] a647fccc = (1024, 4096] d41a514f = (4096, 24064] ee4a9053 = (24064, inf) 99a01f2a * ``sm_107a`` ``bwd_h512``: (0, 1024] 3feed502 = (1024, 4096] 3feed502 = diff --git a/tests/norm/test_cake_rmsnorm_det_train_pairs.py b/tests/norm/test_cake_rmsnorm_det_train_pairs.py index e40a9910426..59ffdf630ff 100644 --- a/tests/norm/test_cake_rmsnorm_det_train_pairs.py +++ b/tests/norm/test_cake_rmsnorm_det_train_pairs.py @@ -247,7 +247,8 @@ def test_rendered_windows_match_the_registry() -> None: """The exporter's ``cake_windows`` module is the registry: its table is ``program_windows`` of every route, its edges are the windows' upper bounds with a known class, and its sections (the docstring notes of the forward and backward) name every BF16 route of their kind and cover the FP32-weight routes in one sentence (they select among - the twins over the same windows), the re-cut ones (``RECUT_FP32_ROUTES``) as lines of their own.""" + the twins over the same windows), the re-cut and swapped ones (``RECUT_FP32_ROUTES``, + ``SWAPPED_FP32_ROUTES``) as lines of their own.""" assert set(windows.PROGRAM_WINDOWS) == set(loader.ROUTES) == set(windows.EDGES) assert set(windows.EDGE_CLASSES) == {"same", "ulp", "dx", "unmeasured"} for arch, routes in loader.ROUTES.items(): @@ -271,9 +272,10 @@ def test_rendered_windows_match_the_registry() -> None: ) listed = f"``{arch}`` ``{route}``:" in windows.SECTIONS[kind] if w_dtype == "f32": - # an FP32-weight route is one sentence of the section, a re-cut one a line of its own; its base - # route is a line either way - recut = (arch, hidden, kind, residual) in RECUT_FP32_ROUTES + # an FP32-weight route is one sentence of the section, a re-cut or swapped one a line of its + # own; its base route is a line either way + key = (arch, hidden, kind, residual) + recut = key in RECUT_FP32_ROUTES or key in SWAPPED_FP32_ROUTES assert listed == recut, (arch, route) base = loader.route_key(hidden, kind, residual) assert f"``{arch}`` ``{base}``:" in windows.SECTIONS[kind], ( @@ -416,6 +418,12 @@ def test_h256_backward_route_runs_the_measured_split(arch: str) -> None: } +# The FP32-weight routes whose window runs the FP32-weight twin of another program than the BF16 route's there, under +# a record of the same host plan (the window, grid rule and workspace of the BF16 record): SM103 H = 1536 above +# 40448 rows, where the register-compensated program is faster with the FP32 weight and not with the BF16 weight. +SWAPPED_FP32_ROUTES = {("sm_103a", 1536, "bwd", False)} + + def _window(record) -> tuple[int | None, int | None]: rule = ( record["select_rule"] or {}