diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_megablocks_cuda_addf474.abi3.so b/build/torch212-cxx11-cu126-x86_64-linux/_megablocks_cuda_addf474.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..5e1ca2387e8e1c62aa97b347d2f8c10bc410fb45 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/_megablocks_cuda_addf474.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:75eac9fccf5aa5c74e862b0aa580f1c420cfc675df9b8b6048ca2b2c3b71ad66 +size 13818088 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so b/build/torch212-cxx11-cu126-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so deleted file mode 100644 index d737d48ad4bbce33084b5d52807a7acf58422021..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7974fda1ad6d9b258e2b469573c2fd106767bdd3d3c25de89efd43b3866e06ad -size 13175016 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_ops.py b/build/torch212-cxx11-cu126-x86_64-linux/_ops.py index 18eab2dbc655830a9f8c1738afc79438a3b0e1f4..6640ac69d26eee39b52fe070e3c8571ed6a93d4e 100644 --- a/build/torch212-cxx11-cu126-x86_64-linux/_ops.py +++ b/build/torch212-cxx11-cu126-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _megablocks_cuda_c4d0dc5 -ops = torch.ops._megablocks_cuda_c4d0dc5 +from . import _megablocks_cuda_addf474 +ops = torch.ops._megablocks_cuda_addf474 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_megablocks_cuda_c4d0dc5::{op_name}" + return f"_megablocks_cuda_addf474::{op_name}" diff --git a/build/torch212-cxx11-cu126-x86_64-linux/megablocks/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/megablocks/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu126-x86_64-linux/megablocks/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/metadata.json b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json index 0577943b41cae1116c6c5095495d3585a486e5e1..32a2b72ac52975a4776496b6c14226907e40ac55 100644 --- a/build/torch212-cxx11-cu126-x86_64-linux/metadata.json +++ b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json @@ -1,9 +1,10 @@ { "name": "megablocks", - "id": "_megablocks_cuda_c4d0dc5", + "id": "_megablocks_cuda_addf474", "version": 1, "license": "Apache-2.0", "python-depends": [], + "kernel-depends": [], "backend": { "type": "cuda", "archs": [ @@ -14,7 +15,7 @@ "8.6", "8.7", "8.9", - "9.0" + "9.0+PTX" ] }, "digest": { @@ -48,8 +49,8 @@ "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=", "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=", "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=", - "_megablocks_cuda_c4d0dc5.abi3.so": "eXT9oa1tmyWOK0aVc8L9EGdnvdPTwl3onv1Ds4ZuBq0=", - "_ops.py": "9NSxkX5lSS5VhF9+hDazP5NimwZK2V/+nIXR0AQxnHA=", + "_megablocks_cuda_addf474.abi3.so": "derJ/M9apcdOhisKpYDxxCDPxnXfm4tgSMorLDtxrWY=", + "_ops.py": "WXlV4XwtKQjdQ6nap3k9TalaHkK0ypeSTshOrV6Hczo=", "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=", "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=", "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=", @@ -61,7 +62,6 @@ "grouped_gemm/ops.py": "Kj7yRgv7afB5SVwXIkN9RNIF2kxD2qrZ0dVsKX5ph2s=", "grouped_gemm_util.py": "bVPrdsEtBVJJ+v/bNrGNf/QHNlVL4VxHg92HvaL3oq4=", "layers.py": "mNTlX2oKJyHBJllRd24lPNdGLa2H9cEUgVLvkZJpoNo=", - "megablocks/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", "ops/__init__.py": "MHQsiMEgkKeZCa8TAYFMGDFnwFIgijfAN7DhMUt/48U=", "ops/all_to_all_benchmark.py": "sxXSCSDdsh+62IWpKNTyBphLTngY1jlgJ1WgNooSEO8=", "ops/binned_gather.py": "pkWne+ThfiTR/bMbDs2QDimKxAI4klqD1FdvhytE7lU=", @@ -106,11 +106,11 @@ "provenance": { "kernel-builder": { "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "commit": "7d2828e6592de641b44fb8eb896719101a2ab2fb", "dirty": false }, "kernel": { - "sha": "c4d0dc56661badb9cc47bde647e08fe3b4bf7458", + "commit": "addf4741fafb55b4c101b15985dfaeea7c7c2f3a", "dirty": false } } diff --git a/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore index 06a6d4e74406eccdedcc14e79124531d98760a95..d9bb8102c26265c6a59a2756b32e390c86ea3fc6 100644 --- a/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore +++ b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore @@ -1 +1 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUTQP9wTaeeOtp+g3scdQtq9C411owCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI1MTIyMTE3WhcNMjYwODI1MTIzMTE3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEPJHV4MhjvJ2wQ8LwG8ho0UIzDoQisDTwmTYSb6o9yaR5vo2Lc5MWf+ps5jmMDZ0koM3SD/zOhuZCPrAI0xoTSaOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUXmifi2X5OX1prd4IqcLgre4p/6kwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGM0ZDBkYzU2NjYxYmFkYjljYzQ3YmRlNjQ3ZTA4ZmUzYjRiZjc0NTgwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzI4NDQ2ODk4OTQvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoDjeIycAAAQDAEcwRQIhAPJGRAotH/fo+g0co0NcZOT6HptOzHfwEx9UqJgnqFr/AiB+rqzjXc5Pjiw5Ks+O4AwTKhC2phb1tJOzdWDO3uQfeDAKBggqhkjOPQQDAwNoADBlAjBEB9Up0jyK6OhoRgTTceucPlZmh+Lb6/GzvJRCN/wotx+NzrsIyp0Z9yF28dKwj8kCMQCvpS9tqg9NLWA6Gwuli8N2tq3DmXC8UmXTuX76YBtVXsL/PtEDz8aOOcmxa+wsZhs="}, "tlogEntries":[{"logIndex":"2583302587", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787660477", "inclusionPromise":{"signedEntryTimestamp":"MEQCIFwqhArAhbAHQvG9biEpGXGnX7Kzy/ag5qMjx3PYHsYpAiAkA7lfjxk0WfIAJ95/WZRQlx6LyhnnP1d9EM/6Na9ONg=="}, "inclusionProof":{"logIndex":"2461398325", "rootHash":"ufLjEJeISMOHLeqKARIA/a15/J4rOiSkz3DMv6ZT5jY=", "treeSize":"2461398326", "hashes":["VRANlHr2KmDIZnL7WzDzhmVmb5oPtDMLaG/OimiIQm4=", "HvIDjbtTaqD74X0cy8sx0UPhrGsYXIpOSqe3GREvdUs=", "tWauNUno2kG7EKvel+T+4ks6XYUxIgWBCrom3zwUgLg=", "BQO6Mvbc6dfqTtcGGvXrh9+QB796m4uHnHENTnE52Ao=", "eUK42vZbtTAA+teDm0LcO+qURyI8IPFvT/lz1znHKQ8=", "r8ZmsIkPd8oW6c3MXr06PoZ4nJYw0Vt5ki3ZDt5yOEM=", "M0t6Pyxtbr61GCmNg60qyyc/qJAEEt98xDnkIgI5BIY=", "EY8T0lxQDrmYoWBGpKNGttaHfKmpfqKcWqSNF4XOQKE=", "T4Jz+m9vO48rOpMSsdZSW28/bgHR8prz7JeNX1s7WiQ=", "EXKrSbeG2qm5APD5QSDN66X7UHz9OSofiW9PyNU7fPM=", "cjIeeutzThxIPTlFFHJqa7bAzn3k2FzK8Nrdd56OWvE=", "+pNXWGIXzU1dqSy/aePHKYtSkDnFinI+/bL5HnurJ6w=", "G1R9F5B2KizFA2NyhDznjNKXPOF1SDmbGHz+omLG7TY=", "RZ7XwGuqMr5gDUh2HpUcaXp+AmIiEcPJiMui4GLJE7o=", "mBm+vQtn0C4thMnlTnxfxqowq1dXsPBCKaJ85da3JeU=", "b80/J/68RC8/tx5xRdlKk6MmTpZiVRotquWBrE/z1pI=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2461398326\nufLjEJeISMOHLeqKARIA/a15/J4rOiSkz3DMv6ZT5jY=\n\n— rekor.sigstore.dev wNI9ajBGAiEAy2+bc6au6mGJM2X55oIkc/qkKalnawOHoulZdgZlEL8CIQCqpCQynrlcGYRMNOA5vrrpBR2IYv/hqoPxykwSYH/9uw==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIyOGRkODcxNzkyODQwNjVhZDgxNDQzZjZlZTFhM2MzNzkyYjI3NWJmNDljYmRlMTkxNGE4YmVkOWNhMzk1ZDM1In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJQkV4M3hNWUdYeXB2c0pYZHhYVEoySUFXQnhuUTc2ZUlHbURqcXZ1cUd4TkFpRUFpVXNMZVNEN2JJS2xIZGlWTEhMZE9mUTZ0b2tjam11aUVjTzg3NHdud1hFPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlZGRlFPWGRVWVdWbFQzUndLMmN6YzJOa1VYUnhPVU0wTVRGdmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1RGTlZFbDVUVlJGTTFkb1kwNU5hbGwzVDBSSk1VMVVTWHBOVkVVelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZRU2toV05FMW9hblpLTW5kUk9FeDNSemhvYnpCVlNYcEViMUZwYzBSVWQyMVVXVk1LWWpadk9YbGhValYyYnpKTVl6Vk5WMllyY0hNMWFtMU5SRm93YTI5Tk0xTkVMM3BQYUhWYVExQnlRVWt3ZUc5VVUyRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZZYldsbUNta3lXRFZQV0RGd2NtUTBTWEZqVEdkeVpUUndMelpyZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVplbEpyVFVkU2FrNVVXVEpPYWtacFdWZFNhVTlYVG1wT1JHUnBDbHBIVlRKT1JHUnNUVVJvYlZwVVRtbE9SMHB0VG5wUk1VOUVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2xhUjFVeVRrUmtiRTFFYUcwS1dsUk9hVTVIU20xT2VsRXhUMFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2tLV2tkVk1rNUVaR3hOUkdodFdsUk9hVTVIU20xT2VsRXhUMFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZE5NRnBFUW1zS1dYcFZNazVxV1hoWmJVWnJXV3BzYWxsNlVUTlpiVkpzVG1wUk0xcFVRVFJhYlZWNldXcFNhVnBxWXpCT1ZHZDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtrMFRrUlJNazlFYXpSUFZGRjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUVhbVZKZVdOQlFVRlJSQXBCUldOM1VsRkphRUZRU2tkU1FXOTBTQzltYnl0bk1HTnZNRTVqV2s5VU5raHdkRTk2U0daM1JYZzVWWEZLWjI1eFJuSXZRV2xDSzNKeGVtcFlZelZRQ21wcGR6VkxjeXRQTkVGM1ZFdG9Rekp3YUdJeGRFcFBlbVJYUkU4emRWRm1aVVJCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcENSVUk1VlhBS01HcDVTelpQYUc5U1oxUlVZMlYxWTFCc1dtMW9LMHhpTmk5SGVuWktVa05PTDNkdmRIZ3JUbnB5YzBsNWNEQmFPWGxHTWpoa1MzZHFPR3REVFZGRGRncHdVemwwY1djNVRreFhRVFpIZDNWc2FUaE9NblJ4TTBSdFdFTTRWVzFZVkhWWU56WlpRblJXV0hOTUwxQjBSVVI2T0dGUFQyTnRlR0VyZDNOYWFITTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg3Zr4gvYHXguQ9HpTQa/zCt0frc5KzBdqeyaaP+z4ig4CFQDt4abH42p4ttb62ZXfRiYymbqw6hgPMjAyNjA4MjUxMjIxMTdaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgyNTEyMjExN1owLwYJKoZIhvcNAQkEMSIEIHn5fQODI/BgWEUrqX3HE5ot0dflrMJBwJfZhlbUV5g1MIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQC9GR6NfulwVg45HhNSg7CnmiCwck7DNLs5vI4vKMtlovk1S+D3qj5FYzWmKh721lECMQD7PEI3N5vc+hrgy86iCZtYEIjZq7k3i0q1/fe3xThkJxFdhFATM9yQJuIdxzsSniA="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"KN2HF5KEBlrYFEP27ho8N5Kydb9Jy94ZFKi+2co5XTU="}, "signature":"MEUCIBEx3xMYGXypvsJXdxXTJ2IAWBxnQ76eIGmDjqvuqGxNAiEAiUsLeSD7bIKlHdiVLHLdOfQ6tokcjmuiEcO874wnwXE="}} \ No newline at end of file +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtCgAwIBAgIUcnefoa2hroWalI1KhhDd0iQTl64wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI3MDcxNDEyWhcNMjYwODI3MDcyNDEyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEqJgLT/s77Ua9ZAQHF1t9p/h16lG+18NSQCJ8w7QoaONmULYsXQB8H0lNJOSvpLSq4fKJ6RHIMhp91fzAsA+9q6OCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUj1e0YQkuCWel+LsKt+YQp0VZbH4wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGFkZGY0NzQxZmFmYjU1YjRjMTAxYjE1OTg1ZGZhZWVhN2M3YzJmM2EwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzMwNDUzOTU5NTIvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoEIRtl8AAAQDAEYwRAIgVgqnPHDnwBnUWG9dr1HfzznpQm7T3O4busbrtbCbcYwCIAjg/mbu6vdd48B3A84zr+2LSc5nnWjmmhGpWGN+MkngMAoGCCqGSM49BAMDA2gAMGUCMFuIJE/JYWzlWth5JuVzDrXTzxniaF+wGxShuHzrTnbF9f7n2HwSNfREY3hnL+GMyQIxAPVCL031YIJ7KtOpjTbaGNG2D6yEZwrwao1NVT9xamWbED9p8sQWk9clGPddIdlihQ=="}, "tlogEntries":[{"logIndex":"2614048455", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787814852", "inclusionPromise":{"signedEntryTimestamp":"MEUCIEnatko5JlMhXJ2NzhOlCrNFOrOvUyfGVkJQ25TCeW6pAiEA+3S3PDF7VtCXDElQZmyzLIdnr17zTR/X5vU6Gro6rnQ="}, "inclusionProof":{"logIndex":"2492144193", "rootHash":"OGbwNrJjc+flOi5K+TnVTDy2pX0o9F8o71DTOhk7M/U=", "treeSize":"2492144221", "hashes":["vhQ4CS2xVVo3D8k32qQJfoKnJ18JfxjlbF4ppD7YiSw=", "LEU31UJM50zZ48YlhexjezNCSAxKcShzRxHsW+6AJlc=", "DDFRUYEZguBgANUZaHKRoXIrcKa3VuV+RTZOFvj3bLk=", "z8MO8JCBH3AJGXDd2kcEeFCSNJAvv8n7F4lZ/yt/Dts=", "BoV8hK3bGHabcdDVPwos21LnIWQXUnA17Axo/WSweSA=", "Rtw5GMrUOIGyS5ueZk0ib5Fs8F3TNSxO4icabC7JhpE=", "gxi1eJ8qX1K1iauxvafmRFx1FhOcSvz3IdnMqR1qs9E=", "VFlTxnDU4W+k/WA3ARTvzCiqVt1yxQIMUyUzMCugJ2Q=", "CFNMzmJ8N+uNx37P2asMqvrW9tw4axAVHnXLRO39tFs=", "NypqAbB+mNix4l1ahGeP9fRWiQzZCm4P1664makjjq4=", "8SoaeA7IeJPoIiveQbBfHCWXSDQabIk2PW+FT49NSEQ=", "VGrFt5JvlIUQvoGEdBBxrCe3lYbu581OhgLq/Xdyat4=", "7TfDCaMWgEwSJEZd/e33b24919rF4IDJLmbqNNRNhZs=", "0HfdxDt/zGugZLIuHasrdLEW9s8OaTykSSqBEDFQXdw=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2492144221\nOGbwNrJjc+flOi5K+TnVTDy2pX0o9F8o71DTOhk7M/U=\n\n— rekor.sigstore.dev wNI9ajBFAiEA4pz/k11O+33XODNFNbBNY0NiRCyXiZJeVMPxwMzqlbkCIDX+thsIWjPlhacQBgZRyQGLTk0ILGRkkkvBH1YsJauI\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI4ZTUxYzRhMjA3NjdkYmIyZmU3OTc1ZDcwNmM1MGM1MjkwM2M0NThkYTBlMzY1ZWE1M2RhYjY3NDlhZjZhNzE0In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRlBheUNPdGNXVHIvWkd5Q2lGNXN0SUwwOTB2MnlJRmRjM0FLWk1aWVhxZEFpRUFpZjdZV2hKWTB6a1YyRVdIZ2tkYXp5RS9OK2NWQUo2RFFXVXcwNzlibmRFPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SRFowRjNTVUpCWjBsVlkyNWxabTloTW1oeWIxZGhiRWt4UzJob1JHUXdhVkZVYkRZMGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1ROTlJHTjRUa1JGZVZkb1kwNU5hbGwzVDBSSk0wMUVZM2xPUkVWNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ4U21kTVZDOXpOemRWWVRsYVFWRklSakYwT1hBdmFERTJiRWNyTVRoT1UxRkRTamdLZHpkUmIyRlBUbTFWVEZseldGRkNPRWd3YkU1S1QxTjJjRXhUY1RSbVMwbzJVa2hKVFdod09URm1la0Z6UVNzNWNUWlBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZxTVdVd0NsbFJhM1ZEVjJWc0sweHpTM1FyV1ZGd01GWmFZa2cwZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVpWMUpyV21wUk0wNUVSbTFaVjFwcFRsUldhVTVIVFhoTlJFWnBDazFVVlRWUFJGWnJXbTFHYkZwWFJUTlplbVJxVFcxWmVsbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWxOVkZVMVQwUldhMXB0Um13S1dsZEZNMWw2WkdwTmJWbDZXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWtLVFZSVk5VOUVWbXRhYlVac1dsZEZNMWw2WkdwTmJWbDZXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZEdhMXBIV1RBS1RucFJlRnB0Um0xWmFsVXhXV3BTYWsxVVFYaFpha1V4VDFSbk1WcEhXbWhhVjFab1RqSk5NMWw2U20xTk1rVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWsxM1RrUlZlazlVVlRWT1ZFbDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUZTVkowYkRoQlFVRlJSQXBCUlZsM1VrRkpaMVpuY1c1UVNFUnVkMEp1VlZkSE9XUnlNVWhtZW5wdWNGRnROMVF6VHpSaWRYTmljblJpUTJKaldYZERTVUZxWnk5dFluVTJkbVJrQ2pRNFFqTkJPRFI2Y2lzeVRGTmpOVzV1VjJwdGJXaEhjRmRIVGl0TmEyNW5UVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlaMEZOUjFWRFRVWjFTVXBGTDBvS1dWZDZiRmQwYURWS2RWWjZSSEpZVkhwNGJtbGhSaXQzUjNoVGFIVkllbkpVYm1KR09XWTNiakpJZDFOT1psSkZXVE5vYmt3clIwMTVVVWw0UVZCV1F3cE1NRE14V1VsS04wdDBUM0JxVkdKaFIwNUhNa1EyZVVWYWQzSjNZVzh4VGxaVU9YaGhiVmRpUlVRNWNEaHpVVmRyT1dOc1IxQmtaRWxrYkdsb1VUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgdy1sI3m2IMSHfw1deoZrxwNhv0qVe+YKzpONjj5oBG4CFDBNEb9YDHm1FD33fH6LxCRTM0E6GA8yMDI2MDgyNzA3MTQxMlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwODI3MDcxNDEyWjAvBgkqhkiG9w0BCQQxIgQgMdZ3f9TM/8ePdAqeloxYU+LMs9Tc4xrwzzu3LvZc1JowgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAIrOlsplb3whBLvuU5JbUVqpvlMe8XeiyxKSmjUgCCWuJk4hDNH6+xi6SMnzMG8l1wIxAOmXiZxrkf8fktq6YAGBuFllN0BuKELXKzFTy6MMgeuJ9tCvOjhttMazMPvyw7HB+g=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"jlHEogdn27L+eXXXBsUMUpA8RY2g42XqU9q2dJr2pxQ="}, "signature":"MEUCIFPayCOtcWTr/ZGyCiF5stIL090v2yIFdc3AKZMZYXqdAiEAif7YWhJY0zkV2EWHgkdazyE/N+cVAJ6DQWUw079bndE="}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_megablocks_cuda_addf474.abi3.so b/build/torch212-cxx11-cu130-x86_64-linux/_megablocks_cuda_addf474.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..b0033523749384f1572239e84883874dafe44e7e --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/_megablocks_cuda_addf474.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b389c414114473ce7b5e8fe84eb3897115c7e63edb6c2536c212356ee37486bd +size 14586872 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so b/build/torch212-cxx11-cu130-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so deleted file mode 100644 index 48bbc18007d92514b4c77cffed12788d87ad55d9..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4d2d01e026bce542464c67904d0306539efdfbd26c866b3839931f531f3be4aa -size 10113176 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_ops.py b/build/torch212-cxx11-cu130-x86_64-linux/_ops.py index 18eab2dbc655830a9f8c1738afc79438a3b0e1f4..6640ac69d26eee39b52fe070e3c8571ed6a93d4e 100644 --- a/build/torch212-cxx11-cu130-x86_64-linux/_ops.py +++ b/build/torch212-cxx11-cu130-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _megablocks_cuda_c4d0dc5 -ops = torch.ops._megablocks_cuda_c4d0dc5 +from . import _megablocks_cuda_addf474 +ops = torch.ops._megablocks_cuda_addf474 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_megablocks_cuda_c4d0dc5::{op_name}" + return f"_megablocks_cuda_addf474::{op_name}" diff --git a/build/torch212-cxx11-cu130-x86_64-linux/megablocks/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/megablocks/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu130-x86_64-linux/megablocks/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/metadata.json b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json index 3bf2050b03c65d153704b5468b32236e82ae4f5c..99efc7e740fe552d886df11c003b3d5a31e12af2 100644 --- a/build/torch212-cxx11-cu130-x86_64-linux/metadata.json +++ b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json @@ -1,14 +1,18 @@ { "name": "megablocks", - "id": "_megablocks_cuda_c4d0dc5", + "id": "_megablocks_cuda_addf474", "version": 1, "license": "Apache-2.0", "python-depends": [], + "kernel-depends": [], "backend": { "type": "cuda", "archs": [ "10.0", + "10.3", + "11.0", "12.0", + "12.1+PTX", "7.5", "8.0", "8.6", @@ -48,8 +52,8 @@ "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=", "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=", "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=", - "_megablocks_cuda_c4d0dc5.abi3.so": "TS0B4Ca85UJGTGeQTQMGU579+9Jshms4OZMfUx875Ko=", - "_ops.py": "9NSxkX5lSS5VhF9+hDazP5NimwZK2V/+nIXR0AQxnHA=", + "_megablocks_cuda_addf474.abi3.so": "s4nEFBFEc857Xo/oTrOJcRXH5j7bbCU2whI1buN0hr0=", + "_ops.py": "WXlV4XwtKQjdQ6nap3k9TalaHkK0ypeSTshOrV6Hczo=", "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=", "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=", "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=", @@ -61,7 +65,6 @@ "grouped_gemm/ops.py": "Kj7yRgv7afB5SVwXIkN9RNIF2kxD2qrZ0dVsKX5ph2s=", "grouped_gemm_util.py": "bVPrdsEtBVJJ+v/bNrGNf/QHNlVL4VxHg92HvaL3oq4=", "layers.py": "mNTlX2oKJyHBJllRd24lPNdGLa2H9cEUgVLvkZJpoNo=", - "megablocks/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", "ops/__init__.py": "MHQsiMEgkKeZCa8TAYFMGDFnwFIgijfAN7DhMUt/48U=", "ops/all_to_all_benchmark.py": "sxXSCSDdsh+62IWpKNTyBphLTngY1jlgJ1WgNooSEO8=", "ops/binned_gather.py": "pkWne+ThfiTR/bMbDs2QDimKxAI4klqD1FdvhytE7lU=", @@ -106,11 +109,11 @@ "provenance": { "kernel-builder": { "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "commit": "7d2828e6592de641b44fb8eb896719101a2ab2fb", "dirty": false }, "kernel": { - "sha": "c4d0dc56661badb9cc47bde647e08fe3b4bf7458", + "commit": "addf4741fafb55b4c101b15985dfaeea7c7c2f3a", "dirty": false } } diff --git a/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore index 8687231b8be9c8464f6c6eda0040ed2b538cae78..e0c60902ea7f4a6b6495ded98997e9de65d920ad 100644 --- a/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore +++ b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore @@ -1 +1 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtCgAwIBAgIUJSd8PABgAeO4DR52EF1GLaff97wwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI1MTIyMTExWhcNMjYwODI1MTIzMTExWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE2J95rQYUscFEvmng+Txabpy9Lj0wMF/PimL1V6P4s2zZpQ3D7BjImDIfhWkbYvHMrhygi05Wkhim+0xt5oouQ6OCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUptKQ2QGeJeH6CPYt8nO+uScnds8wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGM0ZDBkYzU2NjYxYmFkYjljYzQ3YmRlNjQ3ZTA4ZmUzYjRiZjc0NTgwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzI4NDQ2ODk4OTQvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoDjeDocAAAQDAEYwRAIgWfVz6msxL2/OLXywayXxyN3KQNSvgR9bgBEc7ExLXdkCIGNdn4O880wqZp54Q9VpzrJiR/dTGq2y7FR9N/ibNWvdMAoGCCqGSM49BAMDA2kAMGYCMQCoo+jPr8Se4pyJ+0XE95hXL+oyXF0XYu/97P7CoQkVnNz2gEya98dJ/nJF/8uoRzgCMQDd7PPDdEeXiupUhD+BMPTgDzrPFBlVaErHjqXPcKKDanuZwYgfa/NnwLddlkPwOM8="}, "tlogEntries":[{"logIndex":"2583302524", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787660472", "inclusionPromise":{"signedEntryTimestamp":"MEUCID1/KL7+8Sn5TlVyvNAwioEK69P6m5YjZakMUiv5nLZvAiEAlW8nULrUO9h2vPsZwA6T1LKgVXaD2EA3gYREf4Ywp+s="}, "inclusionProof":{"logIndex":"2461398262", "rootHash":"2pXLX5oZTZ/BWtgov4MXYlbWCXrXtoYb3I8T6XBsuig=", "treeSize":"2461398263", "hashes":["mnJetiTJfE1DvJ7lhBRN8OaQX3jLQeNGwI+QLXp0AAM=", "Lh/MtYSxJ8LFybPOQU/g+Frcanyv8Lw29x4fY6wojJc=", "HSg8Qe+EVcgHW9HM8VhTc/rLdK4ZgsiJwWaUtkCDg40=", "dOVbn4NZnSCPsjEt7IgIPnRcYFBIESrnojSV4l++OhI=", "9i5IUFCaI3RtgG9ZfyQpY4loPB/G4SrqefjRx2QlesY=", "s/NH5Bz76m3CS1MP63DS1Iy77fx+9fsdHWF98zIEtG4=", "r8ZmsIkPd8oW6c3MXr06PoZ4nJYw0Vt5ki3ZDt5yOEM=", "M0t6Pyxtbr61GCmNg60qyyc/qJAEEt98xDnkIgI5BIY=", "EY8T0lxQDrmYoWBGpKNGttaHfKmpfqKcWqSNF4XOQKE=", "T4Jz+m9vO48rOpMSsdZSW28/bgHR8prz7JeNX1s7WiQ=", "EXKrSbeG2qm5APD5QSDN66X7UHz9OSofiW9PyNU7fPM=", "cjIeeutzThxIPTlFFHJqa7bAzn3k2FzK8Nrdd56OWvE=", "+pNXWGIXzU1dqSy/aePHKYtSkDnFinI+/bL5HnurJ6w=", "G1R9F5B2KizFA2NyhDznjNKXPOF1SDmbGHz+omLG7TY=", "RZ7XwGuqMr5gDUh2HpUcaXp+AmIiEcPJiMui4GLJE7o=", "mBm+vQtn0C4thMnlTnxfxqowq1dXsPBCKaJ85da3JeU=", "b80/J/68RC8/tx5xRdlKk6MmTpZiVRotquWBrE/z1pI=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2461398263\n2pXLX5oZTZ/BWtgov4MXYlbWCXrXtoYb3I8T6XBsuig=\n\n— rekor.sigstore.dev wNI9ajBEAiAq0xstSQYZOi7qccIojmwFVRzfYnlEejjc7NuwX0OxdwIgAQ4/9zdGDLj4el6CbJOci3IC2YMVwKWxfTHGeZ/05kQ=\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJlN2ZlMzg0OTI4ODY1MTI4ZjNiMTc2NTkwMjI1YzcxMGY3YjVjZDhkNmRhOWFlMWFhODlmN2Y0Y2I1ODUyOGJmIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJQ3RzNHhPcVUvYVdwU2Q5K3lvRXBnamk4ZTY0S0tza3pid3I5ZHBqUkFqa0FpRUE4YmxoUTgwZ0FXMWNubUZKSGJFdS9tTjNOVlh0amdZcTMyUTFuRjNGVjlFPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SRFowRjNTVUpCWjBsVlNsTmtPRkJCUW1kQlpVODBSRkkxTWtWR01VZE1ZV1ptT1RkM2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1RGTlZFbDVUVlJGZUZkb1kwNU5hbGwzVDBSSk1VMVVTWHBOVkVWNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVV5U2prMWNsRlpWWE5qUmtWMmJXNW5LMVI0WVdKd2VUbE1hakIzVFVZdlVHbHRUREVLVmpaUU5ITXllbHB3VVRORU4wSnFTVzFFU1dab1YydGlXWFpJVFhKb2VXZHBNRFZYYTJocGJTc3dlSFExYjI5MVVUWlBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZ3ZEV0UkNqSlJSMlZLWlVnMlExQlpkRGh1VHl0MVUyTnVaSE00ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVplbEpyVFVkU2FrNVVXVEpPYWtacFdWZFNhVTlYVG1wT1JHUnBDbHBIVlRKT1JHUnNUVVJvYlZwVVRtbE9SMHB0VG5wUk1VOUVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2xhUjFVeVRrUmtiRTFFYUcwS1dsUk9hVTVIU20xT2VsRXhUMFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2tLV2tkVk1rNUVaR3hOUkdodFdsUk9hVTVIU20xT2VsRXhUMFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZE5NRnBFUW1zS1dYcFZNazVxV1hoWmJVWnJXV3BzYWxsNlVUTlpiVkpzVG1wUk0xcFVRVFJhYlZWNldXcFNhVnBxWXpCT1ZHZDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtrMFRrUlJNazlFYXpSUFZGRjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUVhbVZFYjJOQlFVRlJSQXBCUlZsM1VrRkpaMWRtVm5vMmJYTjRUREl2VDB4WWVYZGhlVmg0ZVU0elMxRk9VM1puVWpsaVowSkZZemRGZUV4WVpHdERTVWRPWkc0MFR6ZzRNSGR4Q2xwd05UUlJPVlp3ZW5KS2FWSXZaRlJIY1RKNU4wWlNPVTR2YVdKT1YzWmtUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlhMEZOUjFsRFRWRkRiMjhyYWxBS2NqaFRaVFJ3ZVVvck1GaEZPVFZvV0V3cmIzbFlSakJZV1hVdk9UZFFOME52VVd0V2JrNTZNbWRGZVdFNU9HUktMMjVLUmk4NGRXOVNlbWREVFZGRVpBbzNVRkJFWkVWbFdHbDFjRlZvUkN0Q1RWQlVaMFI2Y2xCR1FteFdZVVZ5U0dweFdGQmpTMHRFWVc1MVduZFpaMlpoTDA1dWQweGtaR3hyVUhkUFRUZzlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgbqaIhtv8tLzbp0MAtPYDCkxi33gaCUYaaabJdr6QlKgCFQDt+ch1DN+Vj7pBE+G0P7QGTeilCRgPMjAyNjA4MjUxMjIxMTJaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgyNTEyMjExMlowLwYJKoZIhvcNAQkEMSIEIPV/JPEvcqAyNDTd/+BlkFTb7Y0K9+yO04ncRdnvH9omMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQC+7AsRpZOTkeH34vW6OtA0R/RonZqK8Xjw1a+dGVAo66xxxZZQEuoiwaP3UGXNb/cCMEQZ5CLn5qpxfBtepi1chVLXH8KtmODxOeH7q9tnli4NFxz36Nqm1seveqvAIlo1cA=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"5/44SSiGUSjzsXZZAiXHEPe1zY1tqa4aqJ9/TLWFKL8="}, "signature":"MEUCICts4xOqU/aWpSd9+yoEpgji8e64KKskzbwr9dpjRAjkAiEA8blhQ80gAW1cnmFJHbEu/mN3NVXtjgYq32Q1nF3FV9E="}} \ No newline at end of file +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUA5vul1TIVVVD/55ZQsCMvclTsPUwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI3MDcxNDEzWhcNMjYwODI3MDcyNDEzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEVsmxRDo5AYOQSQTAsN2aOorB4owKn6xe+XSzgWPglS4aVOohNjIFEYxRMkqzqbLIUJofcHvWu8shwCHi4JmWTqOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUsDNkj11Wgu9wXqXUpF8F/w4f0wcwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGFkZGY0NzQxZmFmYjU1YjRjMTAxYjE1OTg1ZGZhZWVhN2M3YzJmM2EwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzMwNDUzOTU5NTIvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoEIRvWoAAAQDAEcwRQIgcR/4wLDiD9qUy4y2vs5kFgcpq3kesrIJjw/f6qC8SuICIQCQNMl1mFo+raHGdMB/QQOkCMRB/d5/fRndBzqLmt7acDAKBggqhkjOPQQDAwNoADBlAjEA/dfyWGtIrwMFzbux13DZsA2hLHUD7kRNRACuN25oBY/2BGqzCB1292Z2oSMLGJa7AjBOiQuqr2faYk9oCfREmjKm8SrbxHWb8WuHAIOgFxGPwADaAHQY/zpoV1v6l5u2blY="}, "tlogEntries":[{"logIndex":"2614048847", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787814854", "inclusionPromise":{"signedEntryTimestamp":"MEQCIHe4p5RPWriaW3uM0phVLA+qZ4hBe+Ywgg82OvTwnOGxAiBuzfj9q2eRheokCj2W33ae37fNGXztKcRriN1wvcjUfg=="}, "inclusionProof":{"logIndex":"2492144585", "rootHash":"GpTg7G2pAAx6ltSe5p5w+mLW5RAZyDa/fxDDYSwnGn0=", "treeSize":"2492144599", "hashes":["Del7EEbYya1uFE6Q2kAS3OZq3gwg3dljdLkc01DaSZY=", "5sL2HZeS6NipYIJ4ciW3HjdMkIB30MyPQQRBPD4VaKI=", "8W861Mu2l0lr8MJykHoaF+w8FUpt2oO2MdJmZMKDGSA=", "bKnOUb0K3WsY+a8UEqANcRu2zO9a82a7lX7Ddi4P21c=", "t+Lp+CPy++8B/TPyzeRJUGruZvhM02wHEpZ6NSeU5FY=", "grNFx/bqJovsvgsk3uMs3EWSG8h43kjyg6PE/8HaWgM=", "izo8eXXW3X4Ups0JGdW/otWtJOUWzCS879a5SRzjnkg=", "SWVtAPePcNx5Tmz8BPDzSaRncMca+mcKIfZKQxBl1/s=", "gxi1eJ8qX1K1iauxvafmRFx1FhOcSvz3IdnMqR1qs9E=", "VFlTxnDU4W+k/WA3ARTvzCiqVt1yxQIMUyUzMCugJ2Q=", "CFNMzmJ8N+uNx37P2asMqvrW9tw4axAVHnXLRO39tFs=", "NypqAbB+mNix4l1ahGeP9fRWiQzZCm4P1664makjjq4=", "8SoaeA7IeJPoIiveQbBfHCWXSDQabIk2PW+FT49NSEQ=", "VGrFt5JvlIUQvoGEdBBxrCe3lYbu581OhgLq/Xdyat4=", "7TfDCaMWgEwSJEZd/e33b24919rF4IDJLmbqNNRNhZs=", "0HfdxDt/zGugZLIuHasrdLEW9s8OaTykSSqBEDFQXdw=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2492144599\nGpTg7G2pAAx6ltSe5p5w+mLW5RAZyDa/fxDDYSwnGn0=\n\n— rekor.sigstore.dev wNI9ajBFAiBxt2lmDypDniaus9sV/HRm7QaAQVOoQtAxiks24mYxvAIhAJBcxCC2pqSDxOdGVGar+cUrIT3Ta9O19cw+35PZR7hQ\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJjOWJlNTllYjMyMWQ2Y2E5ZGE4NzgxMTIxY2RmMjQwMjMzNDU1NTE2NTc3ZDUzMzY0ZmNjODUxNWJkZTVhNDUxIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUM1bS9tZ1FQODk5WjBNMWFPWGQ1NDBwdmY1WmhjZk1zZjVWaGticEhMZkNBSWdXcCtabUdnVEFocG1lQkNaZ0tISWkxUlIwTUVkQmFDOEVWRVpQYUo4S2NJPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlFUVjJkV3d4VkVsV1ZsWkVMelUxV2xGelEwMTJZMnhVYzFCVmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1ROTlJHTjRUa1JGZWxkb1kwNU5hbGwzVDBSSk0wMUVZM2xPUkVWNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZXYzIxNFVrUnZOVUZaVDFGVFVWUkJjMDR5WVU5dmNrSTBiM2RMYmpaNFpTdFlVM29LWjFkUVoyeFROR0ZXVDI5b1RtcEpSa1ZaZUZKTmEzRjZjV0pNU1ZWS2IyWmpTSFpYZFRoemFIZERTR2swU20xWFZIRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZ6UkU1ckNtb3hNVmRuZFRsM1dIRllWWEJHT0VZdmR6Um1NSGRqZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVpWMUpyV21wUk0wNUVSbTFaVjFwcFRsUldhVTVIVFhoTlJFWnBDazFVVlRWUFJGWnJXbTFHYkZwWFJUTlplbVJxVFcxWmVsbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWxOVkZVMVQwUldhMXB0Um13S1dsZEZNMWw2WkdwTmJWbDZXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWtLVFZSVk5VOUVWbXRhYlVac1dsZEZNMWw2WkdwTmJWbDZXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZEdhMXBIV1RBS1RucFJlRnB0Um0xWmFsVXhXV3BTYWsxVVFYaFpha1V4VDFSbk1WcEhXbWhhVjFab1RqSk5NMWw2U20xTk1rVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWsxM1RrUlZlazlVVlRWT1ZFbDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUZTVkoyVjI5QlFVRlJSQXBCUldOM1VsRkpaMk5TTHpSM1RFUnBSRGx4VlhrMGVUSjJjelZyUm1kamNIRXphMlZ6Y2tsS2FuY3Zaalp4UXpoVGRVbERTVkZEVVU1TmJERnRSbThyQ25KaFNFZGtUVUl2VVZGUGEwTk5Va0l2WkRVdlpsSnVaRUo2Y1V4dGREZGhZMFJCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEZRUzlrWm5rS1YwZDBTWEozVFVaNlluVjRNVE5FV25OQk1taE1TRlZFTjJ0U1RsSkJRM1ZPTWpWdlFsa3ZNa0pIY1hwRFFqRXlPVEphTW05VFRVeEhTbUUzUVdwQ1R3cHBVWFZ4Y2pKbVlWbHJPVzlEWmxKRmJXcExiVGhUY21KNFNGZGlPRmQxU0VGSlQyZEdlRWRRZDBGRVlVRklVVmt2ZW5CdlZqRjJObXcxZFRKaWJGazlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgivTPiz46ImUVRh/dqA/JKYgdKFOkokSgBmdjgH2WLIECFFTGQ0mtdpm+MoVyxznY6glXNZ/uGA8yMDI2MDgyNzA3MTQxNFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwODI3MDcxNDE0WjAvBgkqhkiG9w0BCQQxIgQgLKg7YOS45PuElnTBE7vy6grALFrzeI6kw9/l+MXo6P8wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxAN4sMIDl23RF+zZffpzr7zVVJvpcQsFr4/Xr1pG1KcudoghTp531eDpJEshCwwDihAIwXCxgv9HRlvtVb/dQP9pSvOJZ83gnQeIm87q5WPCh9N+BVpOyYuX/rsqINPRHe/2g"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"yb5Z6zIdbKnah4ESHN8kAjNFVRZXfVM2T8yFFb3lpFE="}, "signature":"MEUCIQC5m/mgQP899Z0M1aOXd540pvf5ZhcfMsf5VhkbpHLfCAIgWp+ZmGgTAhpmeBCZgKHIi1RR0MEdBaC8EVEZPaJ8KcI="}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_megablocks_cuda_addf474.abi3.so b/build/torch212-cxx11-cu132-x86_64-linux/_megablocks_cuda_addf474.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..92b2f629fd309f973db2dac67a6f3e34756f3691 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/_megablocks_cuda_addf474.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d6f0251bf07e0c3d494c14363a5b387c25fc59b308c80752270a76ee177a72f5 +size 14986840 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so b/build/torch212-cxx11-cu132-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so deleted file mode 100644 index 3e989a073f9ca829b4004fcca3d2b7a33ecd21b6..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:3cd2b716193592a04cf34b7ae1483011c7ddfcc25489075f354e6cedf1a53a82 -size 10402512 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_ops.py b/build/torch212-cxx11-cu132-x86_64-linux/_ops.py index 18eab2dbc655830a9f8c1738afc79438a3b0e1f4..6640ac69d26eee39b52fe070e3c8571ed6a93d4e 100644 --- a/build/torch212-cxx11-cu132-x86_64-linux/_ops.py +++ b/build/torch212-cxx11-cu132-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _megablocks_cuda_c4d0dc5 -ops = torch.ops._megablocks_cuda_c4d0dc5 +from . import _megablocks_cuda_addf474 +ops = torch.ops._megablocks_cuda_addf474 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_megablocks_cuda_c4d0dc5::{op_name}" + return f"_megablocks_cuda_addf474::{op_name}" diff --git a/build/torch212-cxx11-cu132-x86_64-linux/megablocks/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/megablocks/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch212-cxx11-cu132-x86_64-linux/megablocks/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/metadata.json b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json index dfd2d2c3588d7e9b9aee79c9000732982d98d1be..a87c4f0e764ae4183bc1580b0efeec63a670f195 100644 --- a/build/torch212-cxx11-cu132-x86_64-linux/metadata.json +++ b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json @@ -1,14 +1,18 @@ { "name": "megablocks", - "id": "_megablocks_cuda_c4d0dc5", + "id": "_megablocks_cuda_addf474", "version": 1, "license": "Apache-2.0", "python-depends": [], + "kernel-depends": [], "backend": { "type": "cuda", "archs": [ "10.0", + "10.3", + "11.0", "12.0", + "12.1+PTX", "7.5", "8.0", "8.6", @@ -48,8 +52,8 @@ "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=", "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=", "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=", - "_megablocks_cuda_c4d0dc5.abi3.so": "PNK3Fhk1kqBM80t64UgwEcfd/MJUiQdfNU5s7fGlOoI=", - "_ops.py": "9NSxkX5lSS5VhF9+hDazP5NimwZK2V/+nIXR0AQxnHA=", + "_megablocks_cuda_addf474.abi3.so": "1vAlG/B+DD1JTBQ2Ols4fCX8WbMIyAdSJwp27hd6cvU=", + "_ops.py": "WXlV4XwtKQjdQ6nap3k9TalaHkK0ypeSTshOrV6Hczo=", "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=", "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=", "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=", @@ -61,7 +65,6 @@ "grouped_gemm/ops.py": "Kj7yRgv7afB5SVwXIkN9RNIF2kxD2qrZ0dVsKX5ph2s=", "grouped_gemm_util.py": "bVPrdsEtBVJJ+v/bNrGNf/QHNlVL4VxHg92HvaL3oq4=", "layers.py": "mNTlX2oKJyHBJllRd24lPNdGLa2H9cEUgVLvkZJpoNo=", - "megablocks/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", "ops/__init__.py": "MHQsiMEgkKeZCa8TAYFMGDFnwFIgijfAN7DhMUt/48U=", "ops/all_to_all_benchmark.py": "sxXSCSDdsh+62IWpKNTyBphLTngY1jlgJ1WgNooSEO8=", "ops/binned_gather.py": "pkWne+ThfiTR/bMbDs2QDimKxAI4klqD1FdvhytE7lU=", @@ -106,11 +109,11 @@ "provenance": { "kernel-builder": { "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "commit": "7d2828e6592de641b44fb8eb896719101a2ab2fb", "dirty": false }, "kernel": { - "sha": "c4d0dc56661badb9cc47bde647e08fe3b4bf7458", + "commit": "addf4741fafb55b4c101b15985dfaeea7c7c2f3a", "dirty": false } } diff --git a/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore index bf06897b15fcbf76943796fa57183f57ce068a55..edbae3a2944d5e7ef36645a9e02535f288f37f70 100644 --- a/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore +++ b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore @@ -1 +1 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSTCCBtCgAwIBAgIUOHsK45h2ipRMvVeaYZD9Guc/o+YwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI1MTIyMTEzWhcNMjYwODI1MTIzMTEzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEwhAxJFNYMtNxoPBYzmXY1tZH/5P3ixKgQy1erUTFhE+XzHlDIfF66u3qpKGGN3b3zQTI+seDAopqe4L2F+W6SqOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUYsVomHv+UP1q+6iomqYhMO+nRLIwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGM0ZDBkYzU2NjYxYmFkYjljYzQ3YmRlNjQ3ZTA4ZmUzYjRiZjc0NTgwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzI4NDQ2ODk4OTQvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoDjeFTkAAAQDAEYwRAIgHpfpogDCG6+DnDRUVo3Cvfswjjg85Fn0xaVr1DlVHSECIE5+H/mpf0H51UrS7Gep82qlBLN4t7bY3M5qRv1hvfEbMAoGCCqGSM49BAMDA2cAMGQCMCH9dPE9TAtitu85ZilE+xvqEqQtKVnS1ecYBQrLUvXLmFN6Qsjkx0A7CDjPMsF3WgIwbpeg5+7b7RzyLu8V2hP8VJTynFsiI8lXfY6CJov8ZPdWXk7+BYHshpkwXzIvZxQX"}, "tlogEntries":[{"logIndex":"2583302548", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787660473", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQC81EpDHYqwD0Ft814fa9bgywabWZa1rmGEEbdi6ypdSwIgZS33SeJwbkil3sLdO+EK9OcVYynzONQJxEDc++4blyI="}, "inclusionProof":{"logIndex":"2461398286", "rootHash":"Nd9xPnfnKztd4KuBSM1XAUEh09WalV2yViS27323Qto=", "treeSize":"2461398287", "hashes":["f03VBwyMMrtujA9uh+znFCJ+YJnwPonyHC+iDq1K4h4=", "MVjt7QvYVWoDUKMZYe5kWdnj5kTP6cmEbStyAuhnZ+E=", "HXuvjlv+weMklWPM3GC0X93XS19fBuxDsVxUTRlaKFU=", "eUK42vZbtTAA+teDm0LcO+qURyI8IPFvT/lz1znHKQ8=", "r8ZmsIkPd8oW6c3MXr06PoZ4nJYw0Vt5ki3ZDt5yOEM=", "M0t6Pyxtbr61GCmNg60qyyc/qJAEEt98xDnkIgI5BIY=", "EY8T0lxQDrmYoWBGpKNGttaHfKmpfqKcWqSNF4XOQKE=", "T4Jz+m9vO48rOpMSsdZSW28/bgHR8prz7JeNX1s7WiQ=", "EXKrSbeG2qm5APD5QSDN66X7UHz9OSofiW9PyNU7fPM=", "cjIeeutzThxIPTlFFHJqa7bAzn3k2FzK8Nrdd56OWvE=", "+pNXWGIXzU1dqSy/aePHKYtSkDnFinI+/bL5HnurJ6w=", "G1R9F5B2KizFA2NyhDznjNKXPOF1SDmbGHz+omLG7TY=", "RZ7XwGuqMr5gDUh2HpUcaXp+AmIiEcPJiMui4GLJE7o=", "mBm+vQtn0C4thMnlTnxfxqowq1dXsPBCKaJ85da3JeU=", "b80/J/68RC8/tx5xRdlKk6MmTpZiVRotquWBrE/z1pI=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2461398287\nNd9xPnfnKztd4KuBSM1XAUEh09WalV2yViS27323Qto=\n\n— rekor.sigstore.dev wNI9ajBGAiEAowOyAt2Zb5k4VPUV70DaxxZgHbHEk5IiK8BEzCFFs9oCIQDQbBghR365IJIaOMaHhO1aNufrIrpfrvKbxER0I+6/1g==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJjZGEyYWUxOTE4N2M5MTdhNjczMzE3MjEzMTNiMThmMzVkMGM5ZTZmZmJjNGM4MTI5Zjk5MTczYTA1YWVlMTY1In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJQWtnRStqUXEyL09qU21CZTl6c3Rydm94SEVLSDIzREVkcjg3M0pTU0NuUkFpRUFqV25GTlBacGY0WlBNSE1KVFQ3RHRXTjZjdE1GemRMQTY1ZW8vQVdrMTFvPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRWRU5EUW5SRFowRjNTVUpCWjBsVlQwaHpTelExYURKcGNGSk5kbFpsWVZsYVJEbEhkV012Ynl0WmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1RGTlZFbDVUVlJGZWxkb1kwNU5hbGwzVDBSSk1VMVVTWHBOVkVWNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVYzYUVGNFNrWk9XVTEwVG5odlVFSlplbTFZV1RGMFdrZ3ZOVkF6YVhoTFoxRjVNV1VLY2xWVVJtaEZLMWg2U0d4RVNXWkdOaloxTTNGd1MwZEhUak5pTTNwUlZFa3JjMlZFUVc5d2NXVTBUREpHSzFjMlUzRlBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZaYzFadkNtMUlkaXRWVURGeEt6WnBiMjF4V1doTlR5dHVVa3hKZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVplbEpyVFVkU2FrNVVXVEpPYWtacFdWZFNhVTlYVG1wT1JHUnBDbHBIVlRKT1JHUnNUVVJvYlZwVVRtbE9SMHB0VG5wUk1VOUVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2xhUjFVeVRrUmtiRTFFYUcwS1dsUk9hVTVIU20xT2VsRXhUMFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2tLV2tkVk1rNUVaR3hOUkdodFdsUk9hVTVIU20xT2VsRXhUMFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZE5NRnBFUW1zS1dYcFZNazVxV1hoWmJVWnJXV3BzYWxsNlVUTlpiVkpzVG1wUk0xcFVRVFJhYlZWNldXcFNhVnBxWXpCT1ZHZDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtrMFRrUlJNazlFYXpSUFZGRjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUVhbVZHVkd0QlFVRlJSQXBCUlZsM1VrRkpaMGh3Wm5CdlowUkRSellyUkc1RVVsVldiek5EZG1aemQycHFaemcxUm00d2VHRldjakZFYkZaSVUwVkRTVVUxSzBndmJYQm1NRWcxQ2pGVmNsTTNSMlZ3T0RKeGJFSk1UalIwTjJKWk0wMDFjVkoyTVdoMlprVmlUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlZMEZOUjFGRFRVTklPV1JRUlRrS1ZFRjBhWFIxT0RWYWFXeEZLM2gyY1VWeFVYUkxWbTVUTVdWaldVSlJja3hWZGxoTWJVWk9ObEZ6YW10NE1FRTNRMFJxVUUxelJqTlhaMGwzWW5CbFp3bzFLemRpTjFKNmVVeDFPRll5YUZBNFZrcFVlVzVHYzJsSk9HeFlabGsyUTBwdmRqaGFVR1JYV0dzM0swSlpTSE5vY0d0M1dIcEpkbHA0VVZnS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgIVfsLoyRiTisbRWNoXaZi4aOrfyvcYDrzbHIRZwfkx8CFFkI2OD0ue+GJevJ9BqXxzyQc4/HGA8yMDI2MDgyNTEyMjExM1owAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwODI1MTIyMTEzWjAvBgkqhkiG9w0BCQQxIgQg73jlbewycgDn6ZZtcuFL6G47SQNo1LEi6cvo2v8RA6UwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIwHJK4VKQvu4D8/jY5NxiK+tLnv7SjCF3UvSugLWpdplcPmNXLTawo8u67zCb0r0ofAjEA/HNmVSHIv25cz/JiecPsdW5IvaxbEO9sYtL1+g5CMbKuZSTYb7GKm+YrS7GoFQRg"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"zaKuGRh8kXpnMxchMTsY810Mnm/7xMgSn5kXOgWu4WU="}, "signature":"MEUCIAkgE+jQq2/OjSmBe9zstrvoxHEKH23DEdr873JSSCnRAiEAjWnFNPZpf4ZPMHMJTT7DtWN6ctMFzdLA65eo/AWk11o="}} \ No newline at end of file +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUTZAc00cHnQe/D7dPNb1Q61JdLRcwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI3MDcxNDA3WhcNMjYwODI3MDcyNDA3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAER6YYvzSORzAOvMUg91NjQN/meDWhZESTyVOZvpYArdAJb6XpRxyVKCP2NTiCQb4S67w+SY9aCYbIvU57VJKICKOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUP2uJ7tbXEHAHclfmgpCp2iOIXMgwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGFkZGY0NzQxZmFmYjU1YjRjMTAxYjE1OTg1ZGZhZWVhN2M3YzJmM2EwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzMwNDUzOTU5NTIvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoEIRoyUAAAQDAEcwRQIhANfV92DP21CqpyxEhaIYTKT80B+7DpLmkYIVJ/jqf3tLAiA+sKwiSftQXcOtUSFXjbpIVG/buDUxcIygJrUH9SMGqjAKBggqhkjOPQQDAwNoADBlAjB4+0MVVKnyueeobqse6wpT5yPDhp3HKMbrXdP6ARODY8CURXf/BmrpkLQkMdILBrYCMQDqp0crYZVYeF7P7H/ZjCEWt/laFosGerr4MQafJsxUxnQOvslUyqjcBO9PywyCmhk="}, "tlogEntries":[{"logIndex":"2614047458", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787814847", "inclusionPromise":{"signedEntryTimestamp":"MEQCIF8DaUEUunXnJ+c+NNaVDJYJDPNlO+5ARYoXx/JQxom2AiA1/fS0ErekdpIbpEs+yb8MhRsggClHw/WAGXRFtBoYWg=="}, "inclusionProof":{"logIndex":"2492143196", "rootHash":"elS9QbwVGQqQJHInKI82veZ0COxI5mIm59byPy5PjBw=", "treeSize":"2492143250", "hashes":["3lEfJ3x+2iIiESY9vPI0YvNsiXwoBsw3autJ/n3SGsY=", "WCBemMt4uq8mCZe4bZaHUw1LM+iEc3lo8mAc28L890o=", "ssYcXlLFYUdrfTI+7BhrXFd5V+J7xUklg5BJOdqO4lk=", "UFJAxFXqb/5K0mZZYDspMHmhFDT9m8wcuAOHu0aRqGI=", "R2v+lhN5WjL5xRYlra5RCdfYhV7+9NcKyleW6RMdRTA=", "qETILEsd+PvuF2Ml2mgjSHHcC/s5RjoNDn1C4V/ade8=", "hoqaFa+3T1Dpc6l587gp4kC8JxrVg1UxF1bRVVCC4Dc=", "nFudgbhky7nmBjpXigFw5/wyhomku3S8kQnvmOoP0us=", "74YOSoxcaARV3QmRlp2HcIEFoCumNC8/rDR31M3b3Zk=", "KIag3bHIhHqGqEJl+No8RWATx3rcwFMD/MdDkBYHPmQ=", "CFNMzmJ8N+uNx37P2asMqvrW9tw4axAVHnXLRO39tFs=", "NypqAbB+mNix4l1ahGeP9fRWiQzZCm4P1664makjjq4=", "8SoaeA7IeJPoIiveQbBfHCWXSDQabIk2PW+FT49NSEQ=", "VGrFt5JvlIUQvoGEdBBxrCe3lYbu581OhgLq/Xdyat4=", "7TfDCaMWgEwSJEZd/e33b24919rF4IDJLmbqNNRNhZs=", "0HfdxDt/zGugZLIuHasrdLEW9s8OaTykSSqBEDFQXdw=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2492143250\nelS9QbwVGQqQJHInKI82veZ0COxI5mIm59byPy5PjBw=\n\n— rekor.sigstore.dev wNI9ajBEAiBuapr9w5qEQSdRQyW3vAbWlWgOlHUq6egp0rb3DBhMsAIgdzVnoj2ySLgjrBC5UJzbcoYmQHpRVWG/kWiEliX2rl8=\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJkMDUxZTAwMzQxYzQ5YTY5Mzc5ZTVlNDg4MGE1ODU4Y2ViYmIyZDE4ZmRlYjIxNzdmOWQyZjBhNzMyYzI4M2Y1In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRkN1bmhhZXMwMkxnU3JNREJYNytBb09yYmFFd29IaDYwZTNEa2U0QzMvcUFpRUEvWldiTW82RTZMQTVWbU1MbzYrL3U5cXFsZmlRdXgyVDh2TjJiWVFKSW1VPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlZGcEJZekF3WTBodVVXVXZSRGRrVUU1aU1WRTJNVXBrVEZKamQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1ROTlJHTjRUa1JCTTFkb1kwNU5hbGwzVDBSSk0wMUVZM2xPUkVFelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZTTmxsWmRucFRUMUo2UVU5MlRWVm5PVEZPYWxGT0wyMWxSRmRvV2tWVFZIbFdUMW9LZG5CWlFYSmtRVXBpTmxod1VuaDVWa3REVURKT1ZHbERVV0kwVXpZM2R5dFRXVGxoUTFsaVNYWlZOVGRXU2t0SlEwdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZRTW5WS0NqZDBZbGhGU0VGSVkyeG1iV2R3UTNBeWFVOUpXRTFuZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVpWMUpyV21wUk0wNUVSbTFaVjFwcFRsUldhVTVIVFhoTlJFWnBDazFVVlRWUFJGWnJXbTFHYkZwWFJUTlplbVJxVFcxWmVsbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWxOVkZVMVQwUldhMXB0Um13S1dsZEZNMWw2WkdwTmJWbDZXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWtLVFZSVk5VOUVWbXRhYlVac1dsZEZNMWw2WkdwTmJWbDZXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZEdhMXBIV1RBS1RucFJlRnB0Um0xWmFsVXhXV3BTYWsxVVFYaFpha1V4VDFSbk1WcEhXbWhhVjFab1RqSk5NMWw2U20xTk1rVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWsxM1RrUlZlazlVVlRWT1ZFbDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUZTVkp2ZVZWQlFVRlJSQXBCUldOM1VsRkphRUZPWmxZNU1rUlFNakZEY1hCNWVFVm9ZVWxaVkV0VU9EQkNLemRFY0V4dGExbEpWa292YW5GbU0zUk1RV2xCSzNOTGQybFRablJSQ2xoalQzUlZVMFpZYW1Kd1NWWkhMMkoxUkZWNFkwbDVaMHB5VlVnNVUwMUhjV3BCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcENOQ3N3VFZZS1ZrdHVlWFZsWlc5aWNYTmxObmR3VkRWNVVFUm9jRE5JUzAxaWNsaGtVRFpCVWs5RVdUaERWVkpZWmk5Q2JYSndhMHhSYTAxa1NVeENjbGxEVFZGRWNRcHdNR055V1ZwV1dXVkdOMUEzU0M5YWFrTkZWM1F2YkdGR2IzTkhaWEp5TkUxUllXWktjM2hWZUc1UlQzWnpiRlY1Y1dwalFrODVVSGwzZVVOdGFHczlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgYrdu60lF/YEA9+445V4IwNdNGco3ByUoydVZAZpdHvYCFQCSs1zQcj+u6HVZoauDaeq5o3JgTBgPMjAyNjA4MjcwNzE0MDdaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgyNzA3MTQwN1owLwYJKoZIhvcNAQkEMSIEILeVojiyiC30P6ymBGoC1BDrJXn76/DsnZrjdFCgp6NtMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQCArlEMkaTZxDanNxKF8yn42j/WsWxiPTTPcq3UxP9oVd2s9AJ/LSexVl45gmAYDZ4CMBw0Ld8ZV6RIWwy5lP5g0QaBucSmfUK6JEO2hVy2OM9IHFRqMaBodls1sDJFOMKoQg=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"0FHgA0HEmmk3nl5IgKWFjOu7LRj96yF3+dLwpzLCg/U="}, "signature":"MEUCIFCunhaes02LgSrMDBX7+AoOrbaEwoHh60e3Dke4C3/qAiEA/ZWbMo6E6LA5VmMLo6+/u9qqlfiQux2T8vN2bYQJImU="}} \ No newline at end of file diff --git a/build/torch213-cxx11-cu126-x86_64-linux/_megablocks_cuda_addf474.abi3.so b/build/torch213-cxx11-cu126-x86_64-linux/_megablocks_cuda_addf474.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..81e9b3484fa2f78ddb88edbd12ce62dcd87c965d --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/_megablocks_cuda_addf474.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6069e18a24d5fb1be504d649ce432a8052e1a2bae83896b9723ad581137559bd +size 13822376 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so b/build/torch213-cxx11-cu126-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so deleted file mode 100644 index 42b16b2fd0d7b84dd596bdcdd0d80f756c3d9cfc..0000000000000000000000000000000000000000 --- a/build/torch213-cxx11-cu126-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c2f3017caea39d3305524862086b256082b7936dda2110408241df361cfa47d2 -size 13175208 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/_ops.py b/build/torch213-cxx11-cu126-x86_64-linux/_ops.py index 18eab2dbc655830a9f8c1738afc79438a3b0e1f4..6640ac69d26eee39b52fe070e3c8571ed6a93d4e 100644 --- a/build/torch213-cxx11-cu126-x86_64-linux/_ops.py +++ b/build/torch213-cxx11-cu126-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _megablocks_cuda_c4d0dc5 -ops = torch.ops._megablocks_cuda_c4d0dc5 +from . import _megablocks_cuda_addf474 +ops = torch.ops._megablocks_cuda_addf474 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_megablocks_cuda_c4d0dc5::{op_name}" + return f"_megablocks_cuda_addf474::{op_name}" diff --git a/build/torch213-cxx11-cu126-x86_64-linux/megablocks/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/megablocks/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch213-cxx11-cu126-x86_64-linux/megablocks/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/metadata.json b/build/torch213-cxx11-cu126-x86_64-linux/metadata.json index f809d9bac0aaf5edd85e3f1006fcb8341b1d90c7..00946838ec20a3c0d60247e36a34955e52c18fbf 100644 --- a/build/torch213-cxx11-cu126-x86_64-linux/metadata.json +++ b/build/torch213-cxx11-cu126-x86_64-linux/metadata.json @@ -1,9 +1,10 @@ { "name": "megablocks", - "id": "_megablocks_cuda_c4d0dc5", + "id": "_megablocks_cuda_addf474", "version": 1, "license": "Apache-2.0", "python-depends": [], + "kernel-depends": [], "backend": { "type": "cuda", "archs": [ @@ -14,7 +15,7 @@ "8.6", "8.7", "8.9", - "9.0" + "9.0+PTX" ] }, "digest": { @@ -48,8 +49,8 @@ "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=", "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=", "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=", - "_megablocks_cuda_c4d0dc5.abi3.so": "wvMBfK6jnTMFUkhiCGslYIK3k23aIRBAgkHfNhz6R9I=", - "_ops.py": "9NSxkX5lSS5VhF9+hDazP5NimwZK2V/+nIXR0AQxnHA=", + "_megablocks_cuda_addf474.abi3.so": "YGnhiiTV+xvlBNZJzkMqgFLhorroOJa5cjrVgRN1Wb0=", + "_ops.py": "WXlV4XwtKQjdQ6nap3k9TalaHkK0ypeSTshOrV6Hczo=", "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=", "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=", "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=", @@ -61,7 +62,6 @@ "grouped_gemm/ops.py": "Kj7yRgv7afB5SVwXIkN9RNIF2kxD2qrZ0dVsKX5ph2s=", "grouped_gemm_util.py": "bVPrdsEtBVJJ+v/bNrGNf/QHNlVL4VxHg92HvaL3oq4=", "layers.py": "mNTlX2oKJyHBJllRd24lPNdGLa2H9cEUgVLvkZJpoNo=", - "megablocks/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", "ops/__init__.py": "MHQsiMEgkKeZCa8TAYFMGDFnwFIgijfAN7DhMUt/48U=", "ops/all_to_all_benchmark.py": "sxXSCSDdsh+62IWpKNTyBphLTngY1jlgJ1WgNooSEO8=", "ops/binned_gather.py": "pkWne+ThfiTR/bMbDs2QDimKxAI4klqD1FdvhytE7lU=", @@ -106,11 +106,11 @@ "provenance": { "kernel-builder": { "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "commit": "7d2828e6592de641b44fb8eb896719101a2ab2fb", "dirty": false }, "kernel": { - "sha": "c4d0dc56661badb9cc47bde647e08fe3b4bf7458", + "commit": "addf4741fafb55b4c101b15985dfaeea7c7c2f3a", "dirty": false } } diff --git a/build/torch213-cxx11-cu126-x86_64-linux/metadata.json.sigstore b/build/torch213-cxx11-cu126-x86_64-linux/metadata.json.sigstore index 41b7d50274060ff858255733b87d133e81aac45b..00dea9dc01685b20aed52bb2270bc303ff0cda3a 100644 --- a/build/torch213-cxx11-cu126-x86_64-linux/metadata.json.sigstore +++ b/build/torch213-cxx11-cu126-x86_64-linux/metadata.json.sigstore @@ -1 +1 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtCgAwIBAgIUPKexlqSPtfwNIuqQcSRd2Uk9IoMwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI1MTIyMTE1WhcNMjYwODI1MTIzMTE1WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE4o/8NCXKwvP+gL8EUDN9sggFCqFLCGbzcW4Bo/vCJRQGqWaIklQxd0rjc0yLT3bKeurramzfYL72OH+HMMltMaOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU4bfvGCLTJT31yr+EjuftwCM98xIwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGM0ZDBkYzU2NjYxYmFkYjljYzQ3YmRlNjQ3ZTA4ZmUzYjRiZjc0NTgwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzI4NDQ2ODk4OTQvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoDjeGpgAAAQDAEYwRAIgDQa0Ph1G8Ofv2HreWSKuOSVhSqrSPt6LjRa3HSbEDpcCIFudiw3vzHAun7tCNZg4ceUqfMfrp5tHQqJlEIF6WwWoMAoGCCqGSM49BAMDA2gAMGUCME0BcZCgWkgG0sxuuFOH0t3sFUvYw2GFT/pTBz/Npoh/C0V3UlIi1dvsT2FJwmJrWgIxAO5oAdpf7FkE4Fnf1OuJmrn14JJGEXAp6VprfG1ryrk32S4cTqFm8QzD2iv7EwRJvQ=="}, "tlogEntries":[{"logIndex":"2583302562", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787660475", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQCMWp0psl6m849JTCxD8eYL5YE2PyAUhGC+093vNkJYhgIhANb3ZQBeXt9g2SdEY5e9At+Dm1YeiiwiBJPNwvKatXiI"}, "inclusionProof":{"logIndex":"2461398300", "rootHash":"MBJWZZdAgl5oAXtH7j6jZjGCmgQgxMYKm0nN+UNCmWc=", "treeSize":"2461398304", "hashes":["jlhtoFrVQec0Vb8sv+SpX9+8lOxKnkm6h0z72BSV20U=", "D3IhsQOOBII4ZIpbUfPcqF16+2UwsBKnfyMMsxfelVg=", "L+rX4UOu5EUK9rONsM+nzp0zMzsTFe7vryA0IWQQwqk=", "4QEWTGXLb0zUJxNlZdxkMX7QcytuFcYSeJ1ls224brY=", "dJdYtMdqZPZWIUgGrN0c7iVv8Th0woL+vXNyI9SrD3o=", "eUK42vZbtTAA+teDm0LcO+qURyI8IPFvT/lz1znHKQ8=", "r8ZmsIkPd8oW6c3MXr06PoZ4nJYw0Vt5ki3ZDt5yOEM=", "M0t6Pyxtbr61GCmNg60qyyc/qJAEEt98xDnkIgI5BIY=", "EY8T0lxQDrmYoWBGpKNGttaHfKmpfqKcWqSNF4XOQKE=", "T4Jz+m9vO48rOpMSsdZSW28/bgHR8prz7JeNX1s7WiQ=", "EXKrSbeG2qm5APD5QSDN66X7UHz9OSofiW9PyNU7fPM=", "cjIeeutzThxIPTlFFHJqa7bAzn3k2FzK8Nrdd56OWvE=", "+pNXWGIXzU1dqSy/aePHKYtSkDnFinI+/bL5HnurJ6w=", "G1R9F5B2KizFA2NyhDznjNKXPOF1SDmbGHz+omLG7TY=", "RZ7XwGuqMr5gDUh2HpUcaXp+AmIiEcPJiMui4GLJE7o=", "mBm+vQtn0C4thMnlTnxfxqowq1dXsPBCKaJ85da3JeU=", "b80/J/68RC8/tx5xRdlKk6MmTpZiVRotquWBrE/z1pI=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2461398304\nMBJWZZdAgl5oAXtH7j6jZjGCmgQgxMYKm0nN+UNCmWc=\n\n— rekor.sigstore.dev wNI9ajBFAiEAnvWMyDABFH9IZD0DKx3tw2snxsw1IhgzTWRLVNQNZg4CIDDGFSGfhl369hxfe51XGP4qD6c4Ib7X0S3cUTqhg+Hz\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI0M2Q0NGZiZTg3YzJhOWQwZTJiYjIwZmNlYTI4NzAyNTlkZDE2YzczMTY2NmVhYWQ0NjAzZTBjMjQ2NzQ5NDVkIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUQzSTBPVzBPQ2hyRkVuR0RtMnA5SGhrN25mM1Y1RVJBaXBpbE9pTWtFaEhnSWdBbFhZV0NvdEVpcFFkd2dZVFBIMUJKcDR6VXpWNG9uQzFzN3lhRXRrM3lZPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SRFowRjNTVUpCWjBsVlVFdGxlR3h4VTFCMFpuZE9TWFZ4VVdOVFVtUXlWV3M1U1c5TmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1RGTlZFbDVUVlJGTVZkb1kwNU5hbGwzVDBSSk1VMVVTWHBOVkVVeFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVUwYnk4NFRrTllTM2QyVUN0blREaEZWVVJPT1hOblowWkRjVVpNUTBkaWVtTlhORUlLYnk5MlEwcFNVVWR4VjJGSmEyeFJlR1F3Y21wak1IbE1WRE5pUzJWMWNuSmhiWHBtV1V3M01rOUlLMGhOVFd4MFRXRlBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlUwWW1aMkNrZERURlJLVkRNeGVYSXJSV3AxWm5SM1EwMDVPSGhKZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVplbEpyVFVkU2FrNVVXVEpPYWtacFdWZFNhVTlYVG1wT1JHUnBDbHBIVlRKT1JHUnNUVVJvYlZwVVRtbE9SMHB0VG5wUk1VOUVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2xhUjFVeVRrUmtiRTFFYUcwS1dsUk9hVTVIU20xT2VsRXhUMFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2tLV2tkVk1rNUVaR3hOUkdodFdsUk9hVTVIU20xT2VsRXhUMFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZE5NRnBFUW1zS1dYcFZNazVxV1hoWmJVWnJXV3BzYWxsNlVUTlpiVkpzVG1wUk0xcFVRVFJhYlZWNldXcFNhVnBxWXpCT1ZHZDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtrMFRrUlJNazlFYXpSUFZGRjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUVhbVZIY0dkQlFVRlJSQXBCUlZsM1VrRkpaMFJSWVRCUWFERkhPRTltZGpKSWNtVlhVMHQxVDFOV2FGTnhjbE5RZERaTWFsSmhNMGhUWWtWRWNHTkRTVVoxWkdsM00zWjZTRUYxQ200M2RFTk9XbWMwWTJWVmNXWk5abkp3TlhSSVVYRktiRVZKUmpaWGQxZHZUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlaMEZOUjFWRFRVVXdRbU5hUTJjS1YydG5SekJ6ZUhWMVJrOUlNSFF6YzBaVmRsbDNNa2RHVkM5d1ZFSjZMMDV3YjJndlF6QldNMVZzU1dreFpIWnpWREpHU25kdFNuSlhaMGw0UVU4MWJ3cEJaSEJtTjBaclJUUkdibVl4VDNWS2JYSnVNVFJLU2tkRldFRndObFp3Y21aSE1YSjVjbXN6TWxNMFkxUnhSbTA0VVhwRU1tbDJOMFYzVWtwMlVUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgbmwUpq5yS+9FyOoq5FWeEdmwXYnYqxjOZLCfymM3KP0CFBQEnRhqvI60FFKgzIl+uz/txpqPGA8yMDI2MDgyNTEyMjExNVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwODI1MTIyMTE1WjAvBgkqhkiG9w0BCQQxIgQgVskgydghYfTyg9W7Vmh3NCFVm/2CYRk064SEPT4vuuQwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAKCzAw/Bg+pX2+4Y85xqfcsyHsTgfVPwvne0VtMMXZb0Y8Fv+glXD2MNnH7MPrR8oAIxAOHVQHO/zJaR8mDMi7qSAZH3CwTv4+UEyDdiIv/Gk/AmGzd364ndeEbVJQ9MIcHFfw=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"Q9RPvofCqdDiuyD86ihwJZ3RbHMWZuqtRgPgwkZ0lF0="}, "signature":"MEUCIQD3I0OW0OChrFEnGDm2p9Hhk7nf3V5ERAipilOiMkEhHgIgAlXYWCotEipQdwgYTPH1BJp4zUzV4onC1s7yaEtk3yY="}} \ No newline at end of file +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtGgAwIBAgIUHsSigXA8V+9sG5OzX86VzEeerF0wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI3MDcxNDA5WhcNMjYwODI3MDcyNDA5WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEMhLZCSS9724Q/wP1hrYddGnFsHiNkgWOHI+fgFnFwTeCnuaVIR/vSJaC4+9M7udTneTewkUfnfMDFbkPLMC4Z6OCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUQs0U90i88+tsgNSfVZ9RebGnhJ8wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGFkZGY0NzQxZmFmYjU1YjRjMTAxYjE1OTg1ZGZhZWVhN2M3YzJmM2EwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzMwNDUzOTU5NTIvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoEIRq+YAAAQDAEcwRQIhAOBCYzk9TKa16WftgdDTo5TIJO38ezm5XssKT02txOtBAiBA1EzzIp7Wuq4xPciL1wOwocge+Kyvkx389dXJ0X6RszAKBggqhkjOPQQDAwNnADBkAjAz49tG/pLiwDCm21g9OsHK36n74gBWNgrouMrt4K9HkNKbNWbKbRueP0UM2KZ81OICMBTfSq82r8VzSR69S2bC68Jpvs+D4nTKruV2r+mJ8oM0h3kaJvNB0QPqK/N0GKtm9g=="}, "tlogEntries":[{"logIndex":"2614047887", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787814849", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQC8BmGgNUkm6T/yk1EbJ1NNQ3zrbkJFri7dw3BnlR0e4AIgW85WxEA8YWinXDzfHW25WI/bckkoLtebiAI5L5ivkFI="}, "inclusionProof":{"logIndex":"2492143625", "rootHash":"Z2V6ugp4s7v/3fyWoaVbvtWRKINfIgX1d4DwdnD88YE=", "treeSize":"2492143669", "hashes":["PsU39qYVVAlIp8wI9cxVkqzhdSwCLR+IMX1yb9ZcpJM=", "ePBoPtZHQHI6mGYjaI+Wnq06ShZj6wmXAHA9GiCrr6k=", "4yD1bGiJH93uIYVeh36jqXbecUqJvdMXrsNhhhQq4p4=", "prsuuxeDzdF0aRkSpCTbvJSJ7M1oATAN3hDlk6Okuas=", "KU/ZzA5Z52HbJmBQqm5JPA5I4EERvsuiceOXfei63RU=", "rLjjfxTiWdnJsLWof2wJoZ7bWiVDuGBe8vdD8VYk6jE=", "VFlTxnDU4W+k/WA3ARTvzCiqVt1yxQIMUyUzMCugJ2Q=", "CFNMzmJ8N+uNx37P2asMqvrW9tw4axAVHnXLRO39tFs=", "NypqAbB+mNix4l1ahGeP9fRWiQzZCm4P1664makjjq4=", "8SoaeA7IeJPoIiveQbBfHCWXSDQabIk2PW+FT49NSEQ=", "VGrFt5JvlIUQvoGEdBBxrCe3lYbu581OhgLq/Xdyat4=", "7TfDCaMWgEwSJEZd/e33b24919rF4IDJLmbqNNRNhZs=", "0HfdxDt/zGugZLIuHasrdLEW9s8OaTykSSqBEDFQXdw=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2492143669\nZ2V6ugp4s7v/3fyWoaVbvtWRKINfIgX1d4DwdnD88YE=\n\n— rekor.sigstore.dev wNI9ajBEAiAeemS0veSYb/ZhP0KXnRIewuVB0BCQwhxGI0Be5VIOLQIgWyRQoHecQBAppQqcv3qGcrdUekeQSEzRoH6kGxhTr9A=\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJlNGIyMzhhNTAyNzg5MWU1YWM3YjRkZDRhMTQ2N2Y3MDIwMzY5Y2M1OThmNGJlZDhiOGY5YThhMzc4ZGM4YTdkIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJREd0UEw2cGttUVc2bDlYUVdTbXRmamRZTDA1RitRdW84dWk4OHQvUGNwVkFpRUFvMm1wZUJwOUhYQkdGSEtLRFRFTTBtTXJ3NG5HeDVFeEdmcFhyZnNMcWkwPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SSFowRjNTVUpCWjBsVlNITlRhV2RZUVRoV0t6bHpSelZQZWxnNE5sWjZSV1ZsY2tZd2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1ROTlJHTjRUa1JCTlZkb1kwNU5hbGwzVDBSSk0wMUVZM2xPUkVFMVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZOYUV4YVExTlRPVGN5TkZFdmQxQXhhSEpaWkdSSGJrWnpTR2xPYTJkWFQwaEpLMllLWjBadVJuZFVaVU51ZFdGV1NWSXZkbE5LWVVNMEt6bE5OM1ZrVkc1bFZHVjNhMVZtYm1aTlJFWmlhMUJNVFVNMFdqWlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZSY3pCVkNqa3dhVGc0SzNSelowNVRabFphT1ZKbFlrZHVhRW80ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVpWMUpyV21wUk0wNUVSbTFaVjFwcFRsUldhVTVIVFhoTlJFWnBDazFVVlRWUFJGWnJXbTFHYkZwWFJUTlplbVJxVFcxWmVsbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWxOVkZVMVQwUldhMXB0Um13S1dsZEZNMWw2WkdwTmJWbDZXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWtLVFZSVk5VOUVWbXRhYlVac1dsZEZNMWw2WkdwTmJWbDZXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZEdhMXBIV1RBS1RucFJlRnB0Um0xWmFsVXhXV3BTYWsxVVFYaFpha1V4VDFSbk1WcEhXbWhhVjFab1RqSk5NMWw2U20xTk1rVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWsxM1RrUlZlazlVVlRWT1ZFbDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUZTVkp4SzFsQlFVRlJSQXBCUldOM1VsRkphRUZQUWtOWmVtczVWRXRoTVRaWFpuUm5aRVJVYnpWVVNVcFBNemhsZW0wMVdITnpTMVF3TW5SNFQzUkNRV2xDUVRGRmVucEpjRGRYQ25WeE5IaFFZMmxNTVhkUGQyOWpaMlVyUzNsMmEzZ3pPRGxrV0Vvd1dEWlNjM3BCUzBKblozRm9hMnBQVUZGUlJFRjNUbTVCUkVKclFXcEJlalE1ZEVjS0wzQk1hWGRFUTIweU1XYzVUM05JU3pNMmJqYzBaMEpYVG1keWIzVk5jblEwU3psSWEwNUxZazVYWWt0aVVuVmxVREJWVFRKTFdqZ3hUMGxEVFVKVVpncFRjVGd5Y2poV2VsTlNOamxUTW1KRE5qaEtjSFp6SzBRMGJsUkxjblZXTW5JcmJVbzRiMDB3YUROcllVcDJUa0l3VVZCeFN5OU9NRWRMZEcwNVp6MDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgolMZ3liJt017qRUx/yicdjS/xSe/M030X9M90yGLfbYCFQDz6/FLZLAyE6g0wOg+ZueTePcdixgPMjAyNjA4MjcwNzE0MDlaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgyNzA3MTQwOVowLwYJKoZIhvcNAQkEMSIEINS5YE9VEE5i99un4RvM8i9NvLFVkkAY+OP5b0KJHScVMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMF4AAJ+VQLVtO4m/N+3el75uOvQH/PDfAq82TsURJkMxE4zm2zK+VXXp+jsgcJCTVQIwf1NmqhkhKJx1dwlmF8M5q8d6FwUelOy9j+0af9mFexoxdw2AJ/SmiKCi4UfhLzDB"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"5LI4pQJ4keWse03UoUZ/cCA2nMWY9L7YuPmoo3jcin0="}, "signature":"MEUCIDGtPL6pkmQW6l9XQWSmtfjdYL05F+Quo8ui88t/PcpVAiEAo2mpeBp9HXBGFHKKDTEM0mMrw4nGx5ExGfpXrfsLqi0="}} \ No newline at end of file diff --git a/build/torch213-cxx11-cu130-x86_64-linux/_megablocks_cuda_addf474.abi3.so b/build/torch213-cxx11-cu130-x86_64-linux/_megablocks_cuda_addf474.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..cf76b47be3e42cc32bad07705458cdd2f35c5be0 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/_megablocks_cuda_addf474.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e6cd703144e65044ed60688b698eb888911ff348ec9cc2d205565210431ef06d +size 14587064 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so b/build/torch213-cxx11-cu130-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so deleted file mode 100644 index 92025129c21dc3c102b726eac6189c9567ef5fb7..0000000000000000000000000000000000000000 --- a/build/torch213-cxx11-cu130-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:da57ed156333deb2511dea109ad49184f980f95a5cd13eff0e0a178a817f27ab -size 10113368 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/_ops.py b/build/torch213-cxx11-cu130-x86_64-linux/_ops.py index 18eab2dbc655830a9f8c1738afc79438a3b0e1f4..6640ac69d26eee39b52fe070e3c8571ed6a93d4e 100644 --- a/build/torch213-cxx11-cu130-x86_64-linux/_ops.py +++ b/build/torch213-cxx11-cu130-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _megablocks_cuda_c4d0dc5 -ops = torch.ops._megablocks_cuda_c4d0dc5 +from . import _megablocks_cuda_addf474 +ops = torch.ops._megablocks_cuda_addf474 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_megablocks_cuda_c4d0dc5::{op_name}" + return f"_megablocks_cuda_addf474::{op_name}" diff --git a/build/torch213-cxx11-cu130-x86_64-linux/megablocks/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/megablocks/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch213-cxx11-cu130-x86_64-linux/megablocks/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/metadata.json b/build/torch213-cxx11-cu130-x86_64-linux/metadata.json index e2cd5debf934cc00aadc6a7923ead9f6c73b119d..0ef36e7647bc90e4be6de0520f6c16cf4d8a0eab 100644 --- a/build/torch213-cxx11-cu130-x86_64-linux/metadata.json +++ b/build/torch213-cxx11-cu130-x86_64-linux/metadata.json @@ -1,14 +1,18 @@ { "name": "megablocks", - "id": "_megablocks_cuda_c4d0dc5", + "id": "_megablocks_cuda_addf474", "version": 1, "license": "Apache-2.0", "python-depends": [], + "kernel-depends": [], "backend": { "type": "cuda", "archs": [ "10.0", + "10.3", + "11.0", "12.0", + "12.1+PTX", "7.5", "8.0", "8.6", @@ -48,8 +52,8 @@ "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=", "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=", "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=", - "_megablocks_cuda_c4d0dc5.abi3.so": "2lftFWMz3rJRHeoQmtSRhPmA+Vpc0T7/DgoXioF/J6s=", - "_ops.py": "9NSxkX5lSS5VhF9+hDazP5NimwZK2V/+nIXR0AQxnHA=", + "_megablocks_cuda_addf474.abi3.so": "5s1wMUTmUETtYGiLaY64iJEf80jsnMLSBVZSEEMe8G0=", + "_ops.py": "WXlV4XwtKQjdQ6nap3k9TalaHkK0ypeSTshOrV6Hczo=", "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=", "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=", "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=", @@ -61,7 +65,6 @@ "grouped_gemm/ops.py": "Kj7yRgv7afB5SVwXIkN9RNIF2kxD2qrZ0dVsKX5ph2s=", "grouped_gemm_util.py": "bVPrdsEtBVJJ+v/bNrGNf/QHNlVL4VxHg92HvaL3oq4=", "layers.py": "mNTlX2oKJyHBJllRd24lPNdGLa2H9cEUgVLvkZJpoNo=", - "megablocks/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", "ops/__init__.py": "MHQsiMEgkKeZCa8TAYFMGDFnwFIgijfAN7DhMUt/48U=", "ops/all_to_all_benchmark.py": "sxXSCSDdsh+62IWpKNTyBphLTngY1jlgJ1WgNooSEO8=", "ops/binned_gather.py": "pkWne+ThfiTR/bMbDs2QDimKxAI4klqD1FdvhytE7lU=", @@ -106,11 +109,11 @@ "provenance": { "kernel-builder": { "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "commit": "7d2828e6592de641b44fb8eb896719101a2ab2fb", "dirty": false }, "kernel": { - "sha": "c4d0dc56661badb9cc47bde647e08fe3b4bf7458", + "commit": "addf4741fafb55b4c101b15985dfaeea7c7c2f3a", "dirty": false } } diff --git a/build/torch213-cxx11-cu130-x86_64-linux/metadata.json.sigstore b/build/torch213-cxx11-cu130-x86_64-linux/metadata.json.sigstore index 68f97ffd6255fba419ac816abce307487d3270f1..b687f5fd03adaae5718ca6635da49a7af1b2f051 100644 --- a/build/torch213-cxx11-cu130-x86_64-linux/metadata.json.sigstore +++ b/build/torch213-cxx11-cu130-x86_64-linux/metadata.json.sigstore @@ -1 +1 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHTDCCBtGgAwIBAgIUQ2eE3qLcPa9m8ynfT7fYq7P3T2EwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI1MTIyMTE2WhcNMjYwODI1MTIzMTE2WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEzgKIXs9vLnMRHRNUmzJvBwN30AXyzzEGJdjmRfvgWZkeriOjjDWIWE8uKi2fz54uwjhcGUiefG55pZYCHRBdfaOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUrufvz8pLglIx0xd0HauzenxtyWYwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGM0ZDBkYzU2NjYxYmFkYjljYzQ3YmRlNjQ3ZTA4ZmUzYjRiZjc0NTgwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzI4NDQ2ODk4OTQvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoDjeINwAAAQDAEcwRQIgFM+inYsHBq6IHyTDVTLkzJyzC4mWb1w/X6ZtQWl33HsCIQDe9htRenp1LMRVh9YcYOlRTJKyUbH9A3YnTWfF+h6KsDAKBggqhkjOPQQDAwNpADBmAjEAhRLMp9RgMBeoBjZufmOYUPMktpnPWSMb0VwmEBQRyR1qmqzTO/FSmIpDufLKXtnwAjEA3c59oijmhBLIApFdGgjfhwHvOCxskuKk92bCgwglCvOMZ5FZZtAltWfwGgCGKRKt"}, "tlogEntries":[{"logIndex":"2583302581", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787660476", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQCu5Bw3jdm5PVhJfG1OK/dtP6rFN6XikGK929u60UAJ9QIgdxqksyWuLTskuB5Anbxzj2vv2gKM0twxNewmsHVd1As="}, "inclusionProof":{"logIndex":"2461398319", "rootHash":"tr2Lavb3x9CxvbKytRnxcBRgWo5+GXc540zdYrOiGYU=", "treeSize":"2461398320", "hashes":["5eEexfxfd2NSQvGhcfUn+u2yM9QNOQG1BpmPcpmDWcM=", "uk6tvzglAkgiSfxRJTSKG1TUXds2l1v1uMya5p14L1w=", "AWtcY2WymsIk4LueyDgSQ3fersFRSvrng4fW2anym5c=", "cAJp2HKQxitB324dS3hL9um825VCbKzKrjNKb5AEbZg=", "BQO6Mvbc6dfqTtcGGvXrh9+QB796m4uHnHENTnE52Ao=", "eUK42vZbtTAA+teDm0LcO+qURyI8IPFvT/lz1znHKQ8=", "r8ZmsIkPd8oW6c3MXr06PoZ4nJYw0Vt5ki3ZDt5yOEM=", "M0t6Pyxtbr61GCmNg60qyyc/qJAEEt98xDnkIgI5BIY=", "EY8T0lxQDrmYoWBGpKNGttaHfKmpfqKcWqSNF4XOQKE=", "T4Jz+m9vO48rOpMSsdZSW28/bgHR8prz7JeNX1s7WiQ=", "EXKrSbeG2qm5APD5QSDN66X7UHz9OSofiW9PyNU7fPM=", "cjIeeutzThxIPTlFFHJqa7bAzn3k2FzK8Nrdd56OWvE=", "+pNXWGIXzU1dqSy/aePHKYtSkDnFinI+/bL5HnurJ6w=", "G1R9F5B2KizFA2NyhDznjNKXPOF1SDmbGHz+omLG7TY=", "RZ7XwGuqMr5gDUh2HpUcaXp+AmIiEcPJiMui4GLJE7o=", "mBm+vQtn0C4thMnlTnxfxqowq1dXsPBCKaJ85da3JeU=", "b80/J/68RC8/tx5xRdlKk6MmTpZiVRotquWBrE/z1pI=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2461398320\ntr2Lavb3x9CxvbKytRnxcBRgWo5+GXc540zdYrOiGYU=\n\n— rekor.sigstore.dev wNI9ajBFAiEA4QzUPjnDD+99iJdJiXx8XFQfUXnk4l4I2AhaoqPiA/cCIHzgjuUe64ZA/r5AFnUT7BJabeD0sTDdzQRPYMPYfKEf\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJiMDg2ZTliMTlkNjRhZjIzYzVlYWE3ODk2NjUyNzNhNGI3MTBkMzk5NzljYzNiZmVmNTRmYWI2NzQ2NGZmMDdmIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJRWt5UlA5cUVDc3dNaC9XVjN5NVVjeUJ3RzRuaU5EUVU1dHh6TGJUbjRzS0FpQnRJcWJYbG9KQzVLempOeGU5VTd0bWFpaG1lRjVjVHVNVjhYRGVRZ0pGYnc9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVSRU5EUW5SSFowRjNTVUpCWjBsVlVUSmxSVE54VEdOUVlUbHRPSGx1WmxRM1psbHhOMUF6VkRKRmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1RGTlZFbDVUVlJGTWxkb1kwNU5hbGwzVDBSSk1VMVVTWHBOVkVVeVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVY2WjB0SldITTVka3h1VFZKSVVrNVZiWHBLZGtKM1RqTXdRVmg1ZW5wRlIwcGthbTBLVW1aMloxZGFhMlZ5YVU5cWFrUlhTVmRGT0hWTGFUSm1lalUwZFhkcWFHTkhWV2xsWmtjMU5YQmFXVU5JVWtKa1ptRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZ5ZFdaMkNubzRjRXhuYkVsNE1IaGtNRWhoZFhwbGJuaDBlVmRaZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVplbEpyVFVkU2FrNVVXVEpPYWtacFdWZFNhVTlYVG1wT1JHUnBDbHBIVlRKT1JHUnNUVVJvYlZwVVRtbE9SMHB0VG5wUk1VOUVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2xhUjFVeVRrUmtiRTFFYUcwS1dsUk9hVTVIU20xT2VsRXhUMFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2tLV2tkVk1rNUVaR3hOUkdodFdsUk9hVTVIU20xT2VsRXhUMFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZE5NRnBFUW1zS1dYcFZNazVxV1hoWmJVWnJXV3BzYWxsNlVUTlpiVkpzVG1wUk0xcFVRVFJhYlZWNldXcFNhVnBxWXpCT1ZHZDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtrMFRrUlJNazlFYXpSUFZGRjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUVhbVZKVG5kQlFVRlJSQXBCUldOM1VsRkpaMFpOSzJsdVdYTklRbkUyU1VoNVZFUldWRXhyZWtwNWVrTTBiVmRpTVhjdldEWmFkRkZYYkRNelNITkRTVkZFWlRsb2RGSmxibkF4Q2t4TlVsWm9PVmxqV1U5c1VsUktTM2xWWWtnNVFUTlpibFJYWmtZcmFEWkxjMFJCUzBKblozRm9hMnBQVUZGUlJFRjNUbkJCUkVKdFFXcEZRV2hTVEUwS2NEbFNaMDFDWlc5Q2FscDFabTFQV1ZWUVRXdDBjRzVRVjFOTllqQldkMjFGUWxGU2VWSXhjVzF4ZWxSUEwwWlRiVWx3UkhWbVRFdFlkRzUzUVdwRlFRb3pZelU1YjJscWJXaENURWxCY0Vaa1IyZHFabWgzU0haUFEzaHphM1ZMYXpreVlrTm5kMmRzUTNaUFRWbzFSbHBhZEVGc2RGZG1kMGRuUTBkTFVrdDBDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgoEKiCwnbszWkpyH7aZhp2zp98sIl6+SfdYuchsJCQKMCFCqAIYa0upn0e8tNjfzKIZUnsWheGA8yMDI2MDgyNTEyMjExNlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwODI1MTIyMTE2WjAvBgkqhkiG9w0BCQQxIgQg+cVzcDs30QEUy6HEa5YVdCqI4/EuLnAKSu72arVC5IowgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxAOkiCa9zmmYJqPbwZuCcmeY/rDBjY9Oa+N+l89R9ieJlW7nxNr4XZBejfshbl+xPwwIweWwRFMFw1x9xL2WCxXhppsJDFlytRM6r31+UxHsRxQ4ZKsNUmlj4DAQy3/p0wJSM"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"sIbpsZ1kryPF6qeJZlJzpLcQ05l5zDv+9U+rZ0ZP8H8="}, "signature":"MEQCIEkyRP9qECswMh/WV3y5UcyBwG4niNDQU5txzLbTn4sKAiBtIqbXloJC5KzjNxe9U7tmaihmeF5cTuMV8XDeQgJFbw=="}} \ No newline at end of file +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtGgAwIBAgIUWnYS1zF2bbuoocIpT7LpeSwcVAwwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI3MDcxNDExWhcNMjYwODI3MDcyNDExWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEm56kStkBeKppo36Zy0/O9+ud2mot2lYBInPSztG8Swme3t7G6gR5iZFewIZUMqFpcLFEdz/a7FgHQa8BPkBFHqOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUt64F/s2UHJxlGH/TF4auwMvkOb4wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGFkZGY0NzQxZmFmYjU1YjRjMTAxYjE1OTg1ZGZhZWVhN2M3YzJmM2EwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzMwNDUzOTU5NTIvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoEIRsr4AAAQDAEcwRQIgBd2usNrf/RatExbhKbwl4cPi1tkR3u6HuHBy9dZk6bICIQC7GnHowpAicnoI10ohWCVODenPMLHYwdMV3BghpxfueDAKBggqhkjOPQQDAwNnADBkAjAdeVmivmvK5SCERzLssJT2uc6jn3RMYd5igYOR+G7endAevyFfek2WwwQD1pzJ6+ICMF8fpahkZBn/+gFkiC7qTuSLnF0gGuhDAV6oLkUBqxQNAKbaphG8VZr2mLyxkbkSCA=="}, "tlogEntries":[{"logIndex":"2614048283", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787814851", "inclusionPromise":{"signedEntryTimestamp":"MEUCIEn0dbAX2fh5XVkw9oAnpFNO3GnfeaICyj8JsmYQEdmEAiEAxTfJk5S5kdZXYm09zgjGvKPo6FZQOFyZL2lfd3LndLk="}, "inclusionProof":{"logIndex":"2492144021", "rootHash":"4HEIrictycn6kLT2dBzTnuExmbY1AvWX7mu2EWSzBPY=", "treeSize":"2492144050", "hashes":["03bt7RcBIn3/6pISiUTvoNE6lbQ2UfKmbt2VDhsz1YY=", "YRkfubZKGlRqrQ82R/CF6pMCvRRXVHo/0zYh1OSyTSA=", "cnAywHbi1+mwmV5RxPg+Pfv3PX1BvJyZY0Hei/vdWJw=", "UdCfK0MFItMl/zmZGwFNehWO8i44SumIVuCDPjXT6f8=", "R0VuQNDYNOJty64v4++YKvtG8or5QMzi+v0iWc4y0j4=", "BTWs7zmg6x5lZcg8LJf5XGk6absoEVvidn27LXPsi7c=", "vw4cIKgV5P3+Cz+P5vqtQZc0JkF6TtbPAAowXhOW8KE=", "IlLZ05IZbs5lgJDWCYLdKfpUKQLmCeTy+T6G3jHt+Yc=", "VFlTxnDU4W+k/WA3ARTvzCiqVt1yxQIMUyUzMCugJ2Q=", "CFNMzmJ8N+uNx37P2asMqvrW9tw4axAVHnXLRO39tFs=", "NypqAbB+mNix4l1ahGeP9fRWiQzZCm4P1664makjjq4=", "8SoaeA7IeJPoIiveQbBfHCWXSDQabIk2PW+FT49NSEQ=", "VGrFt5JvlIUQvoGEdBBxrCe3lYbu581OhgLq/Xdyat4=", "7TfDCaMWgEwSJEZd/e33b24919rF4IDJLmbqNNRNhZs=", "0HfdxDt/zGugZLIuHasrdLEW9s8OaTykSSqBEDFQXdw=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2492144050\n4HEIrictycn6kLT2dBzTnuExmbY1AvWX7mu2EWSzBPY=\n\n— rekor.sigstore.dev wNI9ajBFAiEA7Xwme3VhZr+aF8t+9Ei+tWLArAAtQnkh0/s7kf+/T8kCIE8UbaWMTzAEqUbU5E5IYN4jGPBCFNQ9nRpO2p3ZV0Nl\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJmNjEwZGM5YWU2YThmZGMzODRlZTRjZDcwYTlhOThmYjcyYjFkYThjMmJiNjI4Yjc0NGJiYzAzZjg3MmY1OGNlIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJR1BxT1BCeGV6M1pteU5YWkRqN0FSMklPcVBsbHJsVDRndGJaNnpSUDFtL0FpRUE3bXc4WlRTd2tEYUNNRkJML3NXM2lMZ3NNMzF6Q2U2SVdKOVlkRmduZHBBPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SSFowRjNTVUpCWjBsVlYyNVpVekY2UmpKaVluVnZiMk5KY0ZRM1RIQmxVM2RqVmtGM2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1ROTlJHTjRUa1JGZUZkb1kwNU5hbGwzVDBSSk0wMUVZM2xPUkVWNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ0TlRaclUzUnJRbVZMY0hCdk16WmFlVEF2VHprcmRXUXliVzkwTW14WlFrbHVVRk1LZW5SSE9GTjNiV1V6ZERkSE5tZFNOV2xhUm1WM1NWcFZUWEZHY0dOTVJrVmtlaTloTjBablNGRmhPRUpRYTBKR1NIRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlYwTmpSR0NpOXpNbFZJU25oc1IwZ3ZWRVkwWVhWM1RYWnJUMkkwZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVpWMUpyV21wUk0wNUVSbTFaVjFwcFRsUldhVTVIVFhoTlJFWnBDazFVVlRWUFJGWnJXbTFHYkZwWFJUTlplbVJxVFcxWmVsbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWxOVkZVMVQwUldhMXB0Um13S1dsZEZNMWw2WkdwTmJWbDZXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWtLVFZSVk5VOUVWbXRhYlVac1dsZEZNMWw2WkdwTmJWbDZXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZEdhMXBIV1RBS1RucFJlRnB0Um0xWmFsVXhXV3BTYWsxVVFYaFpha1V4VDFSbk1WcEhXbWhhVjFab1RqSk5NMWw2U20xTk1rVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWsxM1RrUlZlazlVVlRWT1ZFbDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUZTVkp6Y2pSQlFVRlJSQXBCUldOM1VsRkpaMEprTW5WelRuSm1MMUpoZEVWNFltaExZbmRzTkdOUWFURjBhMUl6ZFRaSWRVaENlVGxrV21zMllrbERTVkZETjBkdVNHOTNjRUZwQ21OdWIwa3hNRzlvVjBOV1QwUmxibEJOVEVoWmQyUk5Wak5DWjJod2VHWjFaVVJCUzBKblozRm9hMnBQVUZGUlJFRjNUbTVCUkVKclFXcEJaR1ZXYldrS2RtMTJTelZUUTBWU2VreHpjMHBVTW5Wak5tcHVNMUpOV1dRMWFXZFpUMUlyUnpkbGJtUkJaWFo1Um1abGF6SlhkM2RSUkRGd2VrbzJLMGxEVFVZNFpncHdZV2hyV2tKdUx5dG5SbXRwUXpkeFZIVlRURzVHTUdkSGRXaEVRVlkyYjB4clZVSnhlRkZPUVV0aVlYQm9SemhXV25JeWJVeDVlR3RpYTFORFFUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgdIBY9FmMYs0v5LBhLCqPnft0U+RHZP0V8fhzOV1lPBgCFQCpRsmhsENgA9qMO6Zh2nQidUaKlhgPMjAyNjA4MjcwNzE0MTFaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgyNzA3MTQxMVowLwYJKoZIhvcNAQkEMSIEIMlODRHd427Crp8kSiHfaVC73mmpMRNr82e0Qj+M5l5TMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQDKrcjprxpxpA6mZBufEqfOX+aROCy11HnmvAxiTTWsRa1zXfUB8BgdwGlSzoQiQKsCMQDJ8/8g9terKbehxFPdwH9YgJNyUFiTXAjC89qXN3sESjC7WKLnMocGPbtmTrxV4TA="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"9hDcmuao/cOE7kzXCpqY+3Kx2owrtii3RLvAP4cvWM4="}, "signature":"MEUCIGPqOPBxez3ZmyNXZDj7AR2IOqPllrlT4gtbZ6zRP1m/AiEA7mw8ZTSwkDaCMFBL/sW3iLgsM31zCe6IWJ9YdFgndpA="}} \ No newline at end of file diff --git a/build/torch213-cxx11-cu132-x86_64-linux/_megablocks_cuda_addf474.abi3.so b/build/torch213-cxx11-cu132-x86_64-linux/_megablocks_cuda_addf474.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..25dfa6fe15e0b55736142f203b9016aa08f7241b --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/_megablocks_cuda_addf474.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6c590af63c62c28e8ae434d47b9ab628a6e96deeb87b3c1280753fab22101ef4 +size 14987032 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so b/build/torch213-cxx11-cu132-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so deleted file mode 100644 index 6d370f005f79fc5b71060539d8b70a6e8938aa6e..0000000000000000000000000000000000000000 --- a/build/torch213-cxx11-cu132-x86_64-linux/_megablocks_cuda_c4d0dc5.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:884567461acf708fef2b49b8ab9efcf462d0cb8b0751e808386d69b925779146 -size 10402704 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/_ops.py b/build/torch213-cxx11-cu132-x86_64-linux/_ops.py index 18eab2dbc655830a9f8c1738afc79438a3b0e1f4..6640ac69d26eee39b52fe070e3c8571ed6a93d4e 100644 --- a/build/torch213-cxx11-cu132-x86_64-linux/_ops.py +++ b/build/torch213-cxx11-cu132-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _megablocks_cuda_c4d0dc5 -ops = torch.ops._megablocks_cuda_c4d0dc5 +from . import _megablocks_cuda_addf474 +ops = torch.ops._megablocks_cuda_addf474 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_megablocks_cuda_c4d0dc5::{op_name}" + return f"_megablocks_cuda_addf474::{op_name}" diff --git a/build/torch213-cxx11-cu132-x86_64-linux/megablocks/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/megablocks/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch213-cxx11-cu132-x86_64-linux/megablocks/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/metadata.json b/build/torch213-cxx11-cu132-x86_64-linux/metadata.json index 73f33ddc67aea954437c862895c06e8ae9e47dbb..2a11f0ae5bb7474226a1c48dfbb31e72bfa3a92b 100644 --- a/build/torch213-cxx11-cu132-x86_64-linux/metadata.json +++ b/build/torch213-cxx11-cu132-x86_64-linux/metadata.json @@ -1,14 +1,18 @@ { "name": "megablocks", - "id": "_megablocks_cuda_c4d0dc5", + "id": "_megablocks_cuda_addf474", "version": 1, "license": "Apache-2.0", "python-depends": [], + "kernel-depends": [], "backend": { "type": "cuda", "archs": [ "10.0", + "10.3", + "11.0", "12.0", + "12.1+PTX", "7.5", "8.0", "8.6", @@ -48,8 +52,8 @@ "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=", "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=", "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=", - "_megablocks_cuda_c4d0dc5.abi3.so": "iEVnRhrPcI/vK0m4q5789GLQy4sHUegIOG1puSV3kUY=", - "_ops.py": "9NSxkX5lSS5VhF9+hDazP5NimwZK2V/+nIXR0AQxnHA=", + "_megablocks_cuda_addf474.abi3.so": "bFkK9jxiwo6K5DTUe5q2KKbpbe64ezwSgHU/qyIQHvQ=", + "_ops.py": "WXlV4XwtKQjdQ6nap3k9TalaHkK0ypeSTshOrV6Hczo=", "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=", "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=", "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=", @@ -61,7 +65,6 @@ "grouped_gemm/ops.py": "Kj7yRgv7afB5SVwXIkN9RNIF2kxD2qrZ0dVsKX5ph2s=", "grouped_gemm_util.py": "bVPrdsEtBVJJ+v/bNrGNf/QHNlVL4VxHg92HvaL3oq4=", "layers.py": "mNTlX2oKJyHBJllRd24lPNdGLa2H9cEUgVLvkZJpoNo=", - "megablocks/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=", "ops/__init__.py": "MHQsiMEgkKeZCa8TAYFMGDFnwFIgijfAN7DhMUt/48U=", "ops/all_to_all_benchmark.py": "sxXSCSDdsh+62IWpKNTyBphLTngY1jlgJ1WgNooSEO8=", "ops/binned_gather.py": "pkWne+ThfiTR/bMbDs2QDimKxAI4klqD1FdvhytE7lU=", @@ -106,11 +109,11 @@ "provenance": { "kernel-builder": { "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "commit": "7d2828e6592de641b44fb8eb896719101a2ab2fb", "dirty": false }, "kernel": { - "sha": "c4d0dc56661badb9cc47bde647e08fe3b4bf7458", + "commit": "addf4741fafb55b4c101b15985dfaeea7c7c2f3a", "dirty": false } } diff --git a/build/torch213-cxx11-cu132-x86_64-linux/metadata.json.sigstore b/build/torch213-cxx11-cu132-x86_64-linux/metadata.json.sigstore index 917ae1719e3fbea87d7d377b696dd0aca0f07035..97e24572e4bebaabc7ae4bafe30d94be1438f46e 100644 --- a/build/torch213-cxx11-cu132-x86_64-linux/metadata.json.sigstore +++ b/build/torch213-cxx11-cu132-x86_64-linux/metadata.json.sigstore @@ -1 +1 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtGgAwIBAgIUF5DyOMVZVai+Z+yjWIintpmoRvkwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI1MTIyMTE3WhcNMjYwODI1MTIzMTE3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEkHVEgyv/5Z2hPIosGrIjxsDY2SgPvIR3I7ryRAoqPOHxeVnwdMijkh4IOHrrMUTIey5ZtlypNeaGwP9Lop9ds6OCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUc/f8L/9W3JqtmFkrNeFy34wjvBMwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGM0ZDBkYzU2NjYxYmFkYjljYzQ3YmRlNjQ3ZTA4ZmUzYjRiZjc0NTgwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzI4NDQ2ODk4OTQvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoDjeJfYAAAQDAEcwRQIhANvjiRk8H2pUQ1yRN0WIED12UWYHAXymSWQdCA3y+ws/AiAX8lgfooBaUxCltPodrVvBtRYSjEMKn8BdnyXCq4HZWDAKBggqhkjOPQQDAwNnADBkAjA1Tsb8pF1dsG2SidoVomoiipbkqA2Df7LK2NSud1N1i5LvM9Hc+mKpXZfBrTka8+ICMCLzQER1kuNIVc4SiD3L2y1+zW7b2e7oyjbtyZd5eEMjAp8g6HN7H+5QtNu2LiigpQ=="}, "tlogEntries":[{"logIndex":"2583302595", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787660478", "inclusionPromise":{"signedEntryTimestamp":"MEQCID+pplq9/9Kem/px3K+yzWfxjNldVc1fXu6ibh6ufELtAiB7aSn0+iqpy8/wqDa8z+6Ke93A5O9moJSzqkNlZmffLg=="}, "inclusionProof":{"logIndex":"2461398333", "rootHash":"EY2FE8ANXSU2qIj9h4gwrGVFAFad8BQ7eG4+YYD8o5g=", "treeSize":"2461398335", "hashes":["NvMTyutpVc0DFj+m3vEFgR3arquLsDAvYiBtzLizk3w=", "9Rrkg3PlpCmq8cQdSinyYoeLRN0QHQUiXh2sYc8v58k=", "sQkWK7cKmGuj3OiwNkTN/wftR/IJgvbixs2HoJRzOSE=", "gvW/0ym4UwAFu3jSc7Cs5+ZNTAVvHLe1MzxNHz+npCQ=", "tWauNUno2kG7EKvel+T+4ks6XYUxIgWBCrom3zwUgLg=", "BQO6Mvbc6dfqTtcGGvXrh9+QB796m4uHnHENTnE52Ao=", "eUK42vZbtTAA+teDm0LcO+qURyI8IPFvT/lz1znHKQ8=", "r8ZmsIkPd8oW6c3MXr06PoZ4nJYw0Vt5ki3ZDt5yOEM=", "M0t6Pyxtbr61GCmNg60qyyc/qJAEEt98xDnkIgI5BIY=", "EY8T0lxQDrmYoWBGpKNGttaHfKmpfqKcWqSNF4XOQKE=", "T4Jz+m9vO48rOpMSsdZSW28/bgHR8prz7JeNX1s7WiQ=", "EXKrSbeG2qm5APD5QSDN66X7UHz9OSofiW9PyNU7fPM=", "cjIeeutzThxIPTlFFHJqa7bAzn3k2FzK8Nrdd56OWvE=", "+pNXWGIXzU1dqSy/aePHKYtSkDnFinI+/bL5HnurJ6w=", "G1R9F5B2KizFA2NyhDznjNKXPOF1SDmbGHz+omLG7TY=", "RZ7XwGuqMr5gDUh2HpUcaXp+AmIiEcPJiMui4GLJE7o=", "mBm+vQtn0C4thMnlTnxfxqowq1dXsPBCKaJ85da3JeU=", "b80/J/68RC8/tx5xRdlKk6MmTpZiVRotquWBrE/z1pI=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2461398335\nEY2FE8ANXSU2qIj9h4gwrGVFAFad8BQ7eG4+YYD8o5g=\n\n— rekor.sigstore.dev wNI9ajBFAiAa42WA5Z85UvzhojuwnU3XSYXiO3MWSjARjVSwxiyodgIhAK6zi0m6uDXzhG+B0OnBNXYG0V85ll9nsO9fpzHotkvL\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI1MTEzYWE1M2Y5ZmZjZmNlOTBiMjQxNzc1YmMyZTk0NzhlNjgzOTNkZTI1MWM0M2E5YjBkZjYzZmY4NDY4ZTBiIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNFeGZJOE1oVVpVUVAyNm1SUFBHck1jdktTYWJkei9OOTY2UzJHMTB3MHJ3SWdNWnhyRkRDV3JIK3dYQzFSWE9xRmJ4NHlMVXRPNTlmOU1oelkrVXBBck1FPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SSFowRjNTVUpCWjBsVlJqVkVlVTlOVmxwV1lXa3JXaXQ1YWxkSmFXNTBjRzF2VW5acmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1RGTlZFbDVUVlJGTTFkb1kwNU5hbGwzVDBSSk1VMVVTWHBOVkVVelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZyU0ZaRlozbDJMelZhTW1oUVNXOXpSM0pKYW5oelJGa3lVMmRRZGtsU00wazNjbmtLVWtGdmNWQlBTSGhsVm01M1pFMXBhbXRvTkVsUFNISnlUVlZVU1dWNU5WcDBiSGx3VG1WaFIzZFFPVXh2Y0Rsa2N6WlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZqTDJZNENrd3ZPVmN6U25GMGJVWnJjazVsUm5rek5IZHFka0pOZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVplbEpyVFVkU2FrNVVXVEpPYWtacFdWZFNhVTlYVG1wT1JHUnBDbHBIVlRKT1JHUnNUVVJvYlZwVVRtbE9SMHB0VG5wUk1VOUVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2xhUjFVeVRrUmtiRTFFYUcwS1dsUk9hVTVIU20xT2VsRXhUMFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2tLV2tkVk1rNUVaR3hOUkdodFdsUk9hVTVIU20xT2VsRXhUMFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZE5NRnBFUW1zS1dYcFZNazVxV1hoWmJVWnJXV3BzYWxsNlVUTlpiVkpzVG1wUk0xcFVRVFJhYlZWNldXcFNhVnBxWXpCT1ZHZDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtrMFRrUlJNazlFYXpSUFZGRjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUVhbVZLWmxsQlFVRlJSQXBCUldOM1VsRkphRUZPZG1wcFVtczRTREp3VlZFeGVWSk9NRmRKUlVReE1sVlhXVWhCV0hsdFUxZFJaRU5CTTNrcmQzTXZRV2xCV0Roc1oyWnZiMEpoQ2xWNFEyeDBVRzlrY2xaMlFuUlNXVk5xUlUxTGJqaENaRzU1V0VOeE5FaGFWMFJCUzBKblozRm9hMnBQVUZGUlJFRjNUbTVCUkVKclFXcEJNVlJ6WWpnS2NFWXhaSE5ITWxOcFpHOVdiMjF2YVdsd1ltdHhRVEpFWmpkTVN6Sk9VM1ZrTVU0eGFUVk1kazA1U0dNcmJVdHdXRnBtUW5KVWEyRTRLMGxEVFVOTWVncFJSVkl4YTNWT1NWWmpORk5wUkROTU1ua3hLM3BYTjJJeVpUZHZlV3BpZEhsYVpEVmxSVTFxUVhBNFp6WklUamRJS3pWUmRFNTFNa3hwYVdkd1VUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgxeIVPKNrrYtXMZsmfNtCwnPOGdfR9NvP7Njnw/SqP6MCFQDNfmKF4pJ1EFJSHeo/A1ts0Ts+lhgPMjAyNjA4MjUxMjIxMThaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgyNTEyMjExOFowLwYJKoZIhvcNAQkEMSIEIGzt5sr2xY/srtfvIQZ81xvXl5YTCSlFvVU+o9lpfx7gMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQCevyIlSlcjbg7xk5rgexkfxPSSyTmJ3wrlT+JZNwtNaFBSBXkykScNlp6s5eUM3O4CMQDORxw7I/QNIH+FgoN30YR6Hyj4k7S5eGdsAcYa5aO3fPRRgfVoDMIPWbAjr4GWaI4="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"UROqU/n/z86QskF3W8LpR45oOT3iUcQ6mw32P/hGjgs="}, "signature":"MEUCIQCExfI8MhUZUQP26mRPPGrMcvKSabdz/N966S2G10w0rwIgMZxrFDCWrH+wXC1RXOqFbx4yLUtO59f9MhzY+UpArME="}} \ No newline at end of file +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtGgAwIBAgIUa1rY5fShGcQfedCk3dbQrDcT6powCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI3MDcxNDEzWhcNMjYwODI3MDcyNDEzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEng6lK/BeN0Awm0cVKGhiN34YyvwCdQlrFCMUiSm4XwCyjcntR4ZG/r36WpT+YgIt32nKbmIfahanBg7uJrM15KOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUdekMREaT+z0MtF46KfKv7+4mviswHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGFkZGY0NzQxZmFmYjU1YjRjMTAxYjE1OTg1ZGZhZWVhN2M3YzJmM2EwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzMwNDUzOTU5NTIvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoEIRue0AAAQDAEcwRQIhAIqqj4AYFmsf/8LDKF4+6xLDIuF68PEQtTHFDvkGK7i+AiBmobPTu8utL4OC2WTCHCTvZikUlROtWFhagvnFeMrWuzAKBggqhkjOPQQDAwNnADBkAjAmLpAKlJyTYX6DA8NaipyZo8i8e1WN151ZU9n7bga9NWLTTF5W3fQMq75QU4aVXnsCMAWo1SiNV4VwLzu0OxxrRE7pZhXOYRupc7d9ahrZhStaU6ooj5ljFtBVOQ3Dcc/XoA=="}, "tlogEntries":[{"logIndex":"2614048644", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787814853", "inclusionPromise":{"signedEntryTimestamp":"MEUCIHeyRrGFYMLsZ90B4SZ1SP2fZIk5Cri+iPTdCK8wu3syAiEAzGhDlpxrQGZH/4nO/q3/GqXKsAa51gs3/0ahYNPQYRY="}, "inclusionProof":{"logIndex":"2492144382", "rootHash":"515/sFUGPvAeLi6ZM+9gYRAEMYvbTqAOGMePSp4q638=", "treeSize":"2492144419", "hashes":["SZ1QcuHxZpQv12e6F9U5CdtheLQTUnQIZyZ20RW65Ts=", "ClF8WxFR/BX5FNWojg2xA+RSMLCsfFQQ9l0Bz0+UgBI=", "vKY14mV1p23Vw0FaAXydfZlP6nyOMx4zyIi92OOUJAY=", "RzUpHZ++ehuR0OeOKXhJd+5ReeSc7pEJenUy8m/gNQI=", "E0xAnhM1lhenY84RQEwRv6usiyPDP07f2cag248BYSs=", "1mpQ4BfD6S1Tnix5+RB2ABkEzwcXNQmWDrc00zgqZ7g=", "ZAEPc4SjH4mAouX2IuTCEkLhyPlBL0ZdQEHa5DWwDk8=", "/86f/+8yplBih+rHWRuNE/xNx7zHslKJzBzCQZwKfno=", "q0gGRQ2vDBFvDwUIaO6HzxFrMWyjCLgxprzhA3BxcNU=", "gxi1eJ8qX1K1iauxvafmRFx1FhOcSvz3IdnMqR1qs9E=", "VFlTxnDU4W+k/WA3ARTvzCiqVt1yxQIMUyUzMCugJ2Q=", "CFNMzmJ8N+uNx37P2asMqvrW9tw4axAVHnXLRO39tFs=", "NypqAbB+mNix4l1ahGeP9fRWiQzZCm4P1664makjjq4=", "8SoaeA7IeJPoIiveQbBfHCWXSDQabIk2PW+FT49NSEQ=", "VGrFt5JvlIUQvoGEdBBxrCe3lYbu581OhgLq/Xdyat4=", "7TfDCaMWgEwSJEZd/e33b24919rF4IDJLmbqNNRNhZs=", "0HfdxDt/zGugZLIuHasrdLEW9s8OaTykSSqBEDFQXdw=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2492144419\n515/sFUGPvAeLi6ZM+9gYRAEMYvbTqAOGMePSp4q638=\n\n— rekor.sigstore.dev wNI9ajBGAiEA5nM5F5tXv13DDJ6sw6nLDAGUPHNRQ+782WwtX4qJD8gCIQDv3EYlHEWkD+77fTXLyPUEPWkBsgD1PSmFIZT61MqmTA==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI0YzJlNzJlOTU3NjRjNjMxYjk3YThhYTRjN2VlNThlNTU3NWE1NDJkZDhhMjY5MDUwMzE4YmY3ZTgyMTFlNzEzIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJRHpEbW5aN1dYZ3E5bmhOTFBYeGV0NGptZW0vWHhLOVhrRFRoMFhheEJJaEFpQTdxU2t3ZFhSUUROV09IU1BiN1RCaTRLVExCZXNuakVzU09XQ1BRTnkraWc9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SSFowRjNTVUpCWjBsVllURnlXVFZtVTJoSFkxRm1aV1JEYXpOa1lsRnlSR05VTm5CdmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1ROTlJHTjRUa1JGZWxkb1kwNU5hbGwzVDBSSk0wMUVZM2xPUkVWNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ1Wnpac1N5OUNaVTR3UVhkdE1HTldTMGRvYVU0ek5GbDVkbmREWkZGc2NrWkRUVlVLYVZOdE5GaDNRM2xxWTI1MFVqUmFSeTl5TXpaWGNGUXJXV2RKZERNeWJrdGliVWxtWVdoaGJrSm5OM1ZLY2sweE5VdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZrWld0TkNsSkZZVlFyZWpCTmRFWTBOa3RtUzNZM0t6UnRkbWx6ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVpWMUpyV21wUk0wNUVSbTFaVjFwcFRsUldhVTVIVFhoTlJFWnBDazFVVlRWUFJGWnJXbTFHYkZwWFJUTlplbVJxVFcxWmVsbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWxOVkZVMVQwUldhMXB0Um13S1dsZEZNMWw2WkdwTmJWbDZXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWtLVFZSVk5VOUVWbXRhYlVac1dsZEZNMWw2WkdwTmJWbDZXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZEdhMXBIV1RBS1RucFJlRnB0Um0xWmFsVXhXV3BTYWsxVVFYaFpha1V4VDFSbk1WcEhXbWhhVjFab1RqSk5NMWw2U20xTk1rVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWsxM1RrUlZlazlVVlRWT1ZFbDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUZTVkoxWlRCQlFVRlJSQXBCUldOM1VsRkphRUZKY1hGcU5FRlpSbTF6Wmk4NFRFUkxSalFyTm5oTVJFbDFSalk0VUVWUmRGUklSa1IyYTBkTE4ya3JRV2xDYlc5aVVGUjFPSFYwQ2t3MFQwTXlWMVJEU0VOVWRscHBhMVZzVWs5MFYwWm9ZV2QyYmtabFRYSlhkWHBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTVCUkVKclFXcEJiVXh3UVVzS2JFcDVWRmxZTmtSQk9FNWhhWEI1V204NGFUaGxNVmRPTVRVeFdsVTViamRpWjJFNVRsZE1WRlJHTlZjelpsRk5jVGMxVVZVMFlWWllibk5EVFVGWGJ3b3hVMmxPVmpSV2QweDZkVEJQZUhoeVVrVTNjRnBvV0U5WlVuVndZemRrT1dGb2NscG9VM1JoVlRadmIybzFiR3BHZEVKV1QxRXpSR05qTDFodlFUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgvmVnRfBJfjpCzbvskgeP5s87wALCyJooza5/zgrFVjECFQCEXItenUdwVo50Wjw+xiMatUl/ExgPMjAyNjA4MjcwNzE0MTNaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgyNzA3MTQxM1owLwYJKoZIhvcNAQkEMSIEICx6Sjp1NQ6cpS0ZI0oxnWT7hPyku6xwYf86o+4xGnWmMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMETS2x6QVopD2v0Vi6wlWDvPWquJjEzn/khJSRm20btt/mRkhksKm0jiMtqG+Vf/2wIwPq8H1AGa3vIXqYnaP2YWfnNkaU8/y9jT9OVQLP/nNq3Xd9X0HC/Oh/vcLX13YjrB"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"TC5y6VdkxjG5eoqkx+5Y5VdaVC3YomkFAxi/foIR5xM="}, "signature":"MEQCIDzDmnZ7WXgq9nhNLPXxet4jmem/XxK9XkDTh0XaxBIhAiA7qSkwdXRQDNWOHSPb7TBi4KTLBesnjEsSOWCPQNy+ig=="}} \ No newline at end of file diff --git a/build/torch214-cxx11-cu126-x86_64-linux/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0766d7b8da4f97baca212177b4bb989bc6374bf8 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/__init__.py @@ -0,0 +1,205 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch + +# Stable alias: bare `ops` is shadowed by `from . import layers` below. +from ._ops import ops as _compiled_ops +from . import ops + +from .grouped_gemm import backend as gg_backend +from .grouped_gemm import ops as gg_ops + + +from ._layers.arguments import Arguments +from ._layers.dmoe import ParallelDroplessMLP, dMoE +from ._layers.glu import SparseGLU +from ._layers.mlp import MLP, SparseMLP +from ._layers.moe import MoE, ParallelMLP, get_load_balancing_loss + +from . import layers + +# This section contains the direct kernel exports (not inlcuded in the original code) +def exclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor: + """ + Compute exclusive cumulative sum along the specified dimension. + + Args: + x: Input tensor + dim: Dimension along which to compute cumsum + out: Output tensor (modified in-place) + + Returns: + The output tensor + """ + result = ops.exclusive_cumsum(x, dim) + out.copy_(result) + return out + + +def inclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor: + """ + Compute inclusive cumulative sum along the specified dimension. + + Args: + x: Input tensor + dim: Dimension along which to compute cumsum + out: Output tensor (modified in-place) + + Returns: + The output tensor + """ + result = ops.inclusive_cumsum(x, dim) + out.copy_(result) + return out + + +def histogram(x: torch.Tensor, num_bins: int) -> torch.Tensor: + """ + Compute histogram of input tensor values. + + Args: + x: Input tensor + num_bins: Number of histogram bins + + Returns: + Histogram tensor with counts for each bin + """ + return ops.histogram(x, num_bins) + + +def indices( + padded_bins: torch.Tensor, + block_size: int, + output_block_rows: int, + output_block_columns: int, +) -> torch.Tensor: + """ + Construct indices from padded bins for sparse operations. + + Args: + padded_bins: Tensor containing bin boundaries + block_size: Size of each block + output_block_rows: Number of rows in output blocks + output_block_columns: Number of columns in output blocks + + Returns: + Tensor containing constructed indices + """ + return ops.indices(padded_bins, block_size, output_block_rows, output_block_columns) + + +def replicate_forward( + x: torch.Tensor, bins: torch.Tensor, out: torch.Tensor +) -> torch.Tensor: + """ + Forward pass of replicate operation - replicate values according to bin sizes. + + Args: + x: Input tensor with values to replicate + bins: Tensor containing bin sizes + out: Output tensor (modified in-place) + + Returns: + The output tensor + """ + return ops.replicate_forward(x, bins, out) + + +def replicate_backward( + grad: torch.Tensor, bins: torch.Tensor, out: torch.Tensor +) -> torch.Tensor: + """ + Backward pass of replicate operation - reduce gradients back to bins. + + Args: + grad: Gradient tensor to reduce + bins: Tensor containing bin sizes + out: Output tensor (modified in-place) + + Returns: + The output tensor + """ + return ops.replicate_backward(grad, bins, out) + + +def sort( + x: torch.Tensor, end_bit: int, x_out: torch.Tensor, iota_out: torch.Tensor +) -> torch.Tensor: + """ + Radix sort with index tracking. + + Args: + x: Input tensor to sort + end_bit: Number of bits to consider in sorting + x_out: Output tensor for sorted values + iota_out: Output tensor for sorted indices + + Returns: + The sorted values tensor + """ + _compiled_ops.sort(x, end_bit, x_out, iota_out) + return x_out + + +# Convenience functions for common use cases +def cumsum(x: torch.Tensor, dim: int = -1, exclusive: bool = False) -> torch.Tensor: + """ + Compute cumulative sum with automatic output allocation. + + Args: + x: Input tensor + dim: Dimension along which to compute cumsum (default: last dimension) + exclusive: Whether to compute exclusive (True) or inclusive (False) cumsum + + Returns: + New tensor containing the cumulative sum + """ + out = torch.empty_like(x) + if exclusive: + return exclusive_cumsum(x, dim, out) + else: + return inclusive_cumsum(x, dim, out) + + +def argsort(x: torch.Tensor, end_bit: int = 32) -> tuple[torch.Tensor, torch.Tensor]: + """ + Sort tensor and return both sorted values and indices. + + Args: + x: Input tensor to sort + end_bit: Number of bits to consider in sorting + + Returns: + Tuple of (sorted_values, sorted_indices) + """ + x_out = torch.empty_like(x) + iota_out = torch.empty_like(x) + sort(x, end_bit, x_out, iota_out) + return x_out, iota_out + + +# Export public API +__all__ = [ + "MyReplacementLayer", + # Direct kernel exports + "exclusive_cumsum", + "inclusive_cumsum", + "histogram", + "indices", + "replicate_forward", + "replicate_backward", + "sort", + "cumsum", + "argsort", + # Original exports + "Arguments", + "ParallelDroplessMLP", + "dMoE", + "SparseGLU", + "MLP", + "SparseMLP", + "MoE", + "ParallelMLP", + "get_load_balancing_loss", +] diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/_triton_kernels/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/_triton_kernels/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/_triton_kernels/gmm.py b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/_triton_kernels/gmm.py new file mode 100644 index 0000000000000000000000000000000000000000..8c101d07cea416f9390b708e5a35fdc466e48aed --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/_triton_kernels/gmm.py @@ -0,0 +1,574 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. + + +# Imports. +# ------------------------------------------------------------------------------ + +# Python standard library +import functools + +# Triton +import triton +import triton.language as tl + +# AITER +from ..configs import CONFIGS as _CONFIGS +from ..utils._triton import arch_info +from ..utils._triton.pid_preprocessing import pid_grid, remap_xcd + +# Kernel config. +# ------------------------------------------------------------------------------ + + +@functools.lru_cache() +def get_config( + gmm_type: str, M: int, K: int, N: int, G: int, accumulate: bool = False +) -> dict[str, int]: + assert gmm_type in { + "gmm", + "ptgmm", + "nptgmm", + }, f"'{gmm_type}' is an invalid GMM variant." + dev = arch_info.get_arch() + assert ( + dev in _CONFIGS + ), f"No GMM configuration tuned for arch '{dev}'. Supported: {sorted(_CONFIGS)}." + arch_configs = _CONFIGS[dev] + assert ( + "default" in arch_configs[gmm_type] + ), "Default configuration is absent." + key = "accumulate" if accumulate else "default" + return arch_configs[gmm_type][key] + + +# Common code shared by GMM and TGMM kernels. +# ------------------------------------------------------------------------------ + + +# XCD remapping followed by 1D PID to 2D grid mapping. +@triton.jit +def _remap_xcd_tile_grid( + tile_in_mm, + num_row_tiles, + num_col_tiles, + GROUP_SIZE: tl.constexpr = 1, + NUM_XCDS: tl.constexpr = 8, +): + return pid_grid( + remap_xcd(tile_in_mm, num_row_tiles * num_col_tiles, NUM_XCDS=NUM_XCDS), + num_row_tiles, + num_col_tiles, + GROUP_SIZE_M=GROUP_SIZE, + ) + + +# GMM kernel. +# ------------------------------------------------------------------------------ + + +@triton.heuristics( + { + "K_DIVISIBLE_BY_BLOCK_SIZE_K": lambda META: META["K"] % META["BLOCK_SIZE_K"] + == 0, + } +) +@triton.jit +def gmm_kernel( + # Tensor pointers: + lhs_ptr, + rhs_ptr, + group_sizes_ptr, + out_ptr, + bias_ptr, + # Tensor shapes: + M: int, + K: int, + N: int, + G: int, + # Meta-parameters: + TRANS_RHS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + K_DIVISIBLE_BY_BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE: tl.constexpr, + GRID_DIM: tl.constexpr, + USE_BIAS: tl.constexpr, +): + tl.assume(M > 0) + tl.assume(K > 0) + tl.assume(N > 0) + tl.assume(G > 0) + + num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) + tl.device_assert(num_n_tiles > 0, "num_n_tiles <= 0") + + # Current tile. Each program computes multiple tiles of each group. + tile = tl.program_id(0) + tl.device_assert(tile >= 0, "tile < 0 (at initialization)") + + # Tile limit of last MM problem (inclusive). + last_mm_tile = 0 + + # Last input row of lhs and output row of out. Each group reads some rows of + # lhs and writes some rows to out. + last_m = 0 + + # Loop through all (m, K, N) MM problems: + # (m, K) x (K, N) = (m, N) + # sum(m) = M + for g in range(G): + # Get m dimension of current MM problem. + m = tl.load(group_sizes_ptr + g) + # m can be zero if group is empty + tl.device_assert(m >= 0, "m < 0") + + num_m_tiles = tl.cdiv(m, BLOCK_SIZE_M) + # num_m_tiles can be zero if group is empty + tl.device_assert(num_m_tiles >= 0, "num_m_tiles < 0") + + num_tiles = num_m_tiles * num_n_tiles + # num_tiles can be zero if group is empty + tl.device_assert(num_tiles >= 0, "num_tiles < 0") + + # Loop through tiles of current MM problem. + while tile >= last_mm_tile and tile < last_mm_tile + num_tiles: + # Figure out tile coordinates in current MM problem. + tile_in_mm = tile - last_mm_tile + tl.device_assert(tile_in_mm >= 0, "tile_in_mm < 0") + + tile_m, tile_n = _remap_xcd_tile_grid( + tile_in_mm, num_m_tiles, num_n_tiles, GROUP_SIZE=GROUP_SIZE + ) + + # Do regular MM: + + tl.device_assert(tile_m * BLOCK_SIZE_M >= 0, "tile_m * BLOCK_SIZE_M < 0") + tl.device_assert(tile_n * BLOCK_SIZE_N >= 0, "tile_n * BLOCK_SIZE_N < 0") + + offs_lhs_m = ( + tile_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + ) % m + offs_rhs_n = ( + tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + ) % N + offs_k = tl.arange(0, BLOCK_SIZE_K).to(tl.int64) + + lhs_ptrs = lhs_ptr + (last_m + offs_lhs_m[:, None]) * K + offs_k[None, :] + + if TRANS_RHS: + rhs_ptrs = ( + rhs_ptr + + g.to(tl.int64) * K * N + + offs_k[:, None] + + offs_rhs_n[None, :] * K + ) + else: + rhs_ptrs = ( + rhs_ptr + + g.to(tl.int64) * K * N + + offs_k[:, None] * N + + offs_rhs_n[None, :] + ) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + if K_DIVISIBLE_BY_BLOCK_SIZE_K: + lhs = tl.load(lhs_ptrs) + rhs = tl.load(rhs_ptrs) + else: + k_mask_limit = K - k * BLOCK_SIZE_K + lhs = tl.load( + lhs_ptrs, mask=offs_k[None, :] < k_mask_limit, other=0 + ) + rhs = tl.load( + rhs_ptrs, mask=offs_k[:, None] < k_mask_limit, other=0 + ) + + acc = tl.dot(lhs, rhs, acc=acc) + + lhs_ptrs += BLOCK_SIZE_K + + if TRANS_RHS: + rhs_ptrs += BLOCK_SIZE_K + else: + rhs_ptrs += BLOCK_SIZE_K * N + + # Add bias if enabled + if USE_BIAS: + offs_bias_n = tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange( + 0, BLOCK_SIZE_N + ) + bias_ptrs = bias_ptr + g.to(tl.int64) * N + offs_bias_n + bias = tl.load(bias_ptrs, mask=offs_bias_n < N, other=0.0) + # Convert bias to float32 to match accumulator precision + bias = bias.to(tl.float32) + # Broadcast bias across M dimension and add in float32 + acc += bias[None, :] + + # Convert to output dtype after all computations + acc = acc.to(out_ptr.type.element_ty) + + offs_out_m = tile_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_out_n = tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + + out_ptrs = ( + out_ptr + (last_m + offs_out_m[:, None]) * N + offs_out_n[None, :] + ) + + tl.store( + out_ptrs, + acc, + mask=(offs_out_m[:, None] < m) & (offs_out_n[None, :] < N), + ) + + # Go to the next tile by advancing number of programs. + tile += GRID_DIM + tl.device_assert(tile > 0, "tile <= 0 (at update)") + + # Get ready to go to the next MM problem. + + last_mm_tile += num_tiles + # last_mm_tile can be zero if group 0 is skipped + tl.device_assert(last_mm_tile >= 0, "last_mm_tile < 0 (at update)") + + last_m += m + # last_m can be zero if group 0 is skipped + tl.device_assert(last_m >= 0, "last_m < 0 (at update)") + tl.device_assert(last_m <= M, "last_m > M (at update)") + + +# Persistent TGMM kernel. +# ------------------------------------------------------------------------------ + + +@triton.jit +def tgmm_persistent_kernel( + # Tensor pointers: + lhs_ptr, + rhs_ptr, + group_sizes_ptr, + out_ptr, + bias_grad_ptr, + # Tensor shapes: + M: int, + K: int, + N: int, + G: int, + # Meta-parameters: + TRANS_LHS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + GROUP_SIZE: tl.constexpr, + GRID_DIM: tl.constexpr, + COMPUTE_BIAS_GRAD: tl.constexpr, + ACCUMULATE: tl.constexpr, +): + tl.assume(M > 0) + tl.assume(K > 0) + tl.assume(N > 0) + tl.assume(G > 0) + + num_k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + tl.device_assert(num_k_tiles > 0, "num_k_tiles <= 0") + + num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) + tl.device_assert(num_n_tiles > 0, "num_n_tiles <= 0") + + num_tiles = num_k_tiles * num_n_tiles + tl.device_assert(num_tiles > 0, "num_tiles <= 0") + + # Current tile. Each program computes multiple tiles of each group. + tile = tl.program_id(0) + tl.device_assert(tile >= 0, "tile < 0 (at initialization)") + + # Tile limit of last MM problem (inclusive). + last_mm_tile = 0 + + # Last input column of lhs and input row of rhs. Each group reads some + # columns of lhs and some rows of rhs. + last_m = 0 + + # Loop through all (K, m, N) MM problems: + # (K, m) x (m, N) = (K, N) + # sum(m) = M + for g in range(G): + # Get m dimension of current MM problem. + m = tl.load(group_sizes_ptr + g) + # m can be zero if group is empty + tl.device_assert(m >= 0, "m < 0") + + # Loop through tiles of current MM problem. + while tile >= last_mm_tile and tile < last_mm_tile + num_tiles: + # Figure out tile coordinates in current MM problem. + tile_in_mm = tile - last_mm_tile + tl.device_assert(tile_in_mm >= 0, "tile_in_mm < 0") + + tile_k, tile_n = _remap_xcd_tile_grid( + tile_in_mm, num_k_tiles, num_n_tiles, GROUP_SIZE=GROUP_SIZE + ) + + # Do regular MM: + + tl.device_assert(tile_k * BLOCK_SIZE_K >= 0, "tile_k * BLOCK_SIZE_K < 0") + tl.device_assert(tile_n * BLOCK_SIZE_N >= 0, "tile_n * BLOCK_SIZE_N < 0") + + offs_lhs_k = ( + tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + ) % K + offs_rhs_n = ( + tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + ) % N + offs_m = tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + + if TRANS_LHS: + lhs_ptrs = ( + lhs_ptr + offs_lhs_k[:, None] + (last_m + offs_m[None, :]) * K + ) + else: + lhs_ptrs = ( + lhs_ptr + offs_lhs_k[:, None] * M + (last_m + offs_m[None, :]) + ) + + rhs_ptrs = rhs_ptr + (last_m + offs_m[:, None]) * N + offs_rhs_n[None, :] + + loop_m = tl.cdiv(m, BLOCK_SIZE_M) + m_divisible_by_block_m = m % BLOCK_SIZE_M == 0 + if not m_divisible_by_block_m: + loop_m -= 1 + + acc = tl.zeros((BLOCK_SIZE_K, BLOCK_SIZE_N), dtype=tl.float32) + + # Initialize bias accumulator + bias_acc = tl.zeros((BLOCK_SIZE_K,), dtype=tl.float32) + + for _ in range(0, loop_m): + lhs = tl.load(lhs_ptrs) + rhs = tl.load(rhs_ptrs) + + acc = tl.dot(lhs, rhs, acc=acc) + + # Accumulate for bias gradient: sum lhs across M dimension + if COMPUTE_BIAS_GRAD and tile_n == 0: + bias_acc += tl.sum( + lhs, axis=1 + ) # Sum across M dimension [K, M] -> [K] + + if TRANS_LHS: + lhs_ptrs += BLOCK_SIZE_M * K + else: + lhs_ptrs += BLOCK_SIZE_M + + rhs_ptrs += BLOCK_SIZE_M * N + + if not m_divisible_by_block_m: + offs_lhs_k = ( + tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + ) % K + offs_rhs_n = ( + tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + ) % N + offs_m = loop_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + lhs = tl.load(lhs_ptrs, mask=offs_m[None, :] < m, other=0) + rhs = tl.load(rhs_ptrs, mask=offs_m[:, None] < m, other=0) + acc = tl.dot(lhs, rhs, acc=acc) + + # Accumulate last chunk for bias gradient + if COMPUTE_BIAS_GRAD and tile_n == 0: + bias_acc += tl.sum(lhs, axis=1) + + acc = acc.to(out_ptr.type.element_ty) + + offs_out_k = tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + offs_out_n = tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + + out_ptrs = ( + out_ptr + + g.to(tl.int64) * K * N + + offs_out_k[:, None] * N + + offs_out_n[None, :] + ) + + mask = (offs_out_k[:, None] < K) & (offs_out_n[None, :] < N) + if ACCUMULATE: + # Load existing values and add to them (like beta=1 in BLAS) + old_vals = tl.load(out_ptrs, mask=mask, other=0.0) + tl.store(out_ptrs, acc + old_vals, mask=mask) + else: + # Overwrite output (like beta=0 in BLAS) + tl.store(out_ptrs, acc, mask=mask) + + # Store bias gradient (only for first N tile, sum across all M) + if COMPUTE_BIAS_GRAD and tile_n == 0: + # Keep as float32 for atomic_add (bf16 not supported for atomics) + bias_grad_ptrs = bias_grad_ptr + g.to(tl.int64) * K + offs_out_k + # Use atomic add since multiple K-tiles may write to same expert's bias + tl.atomic_add( + bias_grad_ptrs, bias_acc, mask=offs_out_k < K, sem="relaxed" + ) + + # Go to the next tile by advancing number of programs. + tile += GRID_DIM + tl.device_assert(tile > 0, "tile <= 0 (at update)") + + # Get ready to go to the next MM problem. + + last_mm_tile += num_tiles + # last_mm_tile can be zero if group 0 is skipped + tl.device_assert(last_mm_tile >= 0, "last_mm_tile < 0 (at update)") + + last_m += m + # last_m can be zero if group 0 is skipped + tl.device_assert(last_m >= 0, "last_m < 0 (at update)") + tl.device_assert(last_m <= M, "last_m > M (at update)") + + +# Regular non-persistent TGMM kernel. +# ------------------------------------------------------------------------------ + + +@triton.heuristics({"BLOCK_SIZE_G": lambda META: triton.next_power_of_2(META["G"])}) +@triton.jit +def tgmm_non_persistent_kernel( + # Tensor pointers: + lhs_ptr, + rhs_ptr, + group_sizes_ptr, + out_ptr, + bias_grad_ptr, + # Tensor shapes: + M: int, + K: int, + N: int, + G: int, + # Meta-parameters: + TRANS_LHS: tl.constexpr, + BLOCK_SIZE_G: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + GROUP_SIZE: tl.constexpr, + COMPUTE_BIAS_GRAD: tl.constexpr, + ACCUMULATE: tl.constexpr, +): + tl.assume(M > 0) + tl.assume(K > 0) + tl.assume(N > 0) + tl.assume(G > 0) + + # Get group ID from grid. + g = tl.program_id(0) + tl.device_assert(g >= 0, "g < 0") + tl.device_assert(g < G, "g >= G") + + # Get m dimension of current MM group. + m = tl.load(group_sizes_ptr + g) + # m can be zero if group is empty. + tl.device_assert(m >= 0, "m < 0") + + # Skip empty groups. + if m == 0: + return + + # Compute sum(group_sizes) until current group g. + # It's the starting column of lhs and starting row of rhs. + offs_g = tl.arange(0, BLOCK_SIZE_G) + group_sizes = tl.load(group_sizes_ptr + offs_g, mask=offs_g < g, other=0) + start_m = tl.sum(group_sizes) + + num_k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + tl.device_assert(num_k_tiles > 0, "num_k_tiles <= 0") + + num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) + tl.device_assert(num_n_tiles > 0, "num_n_tiles <= 0") + + # Get MM tile from grid. + tile_in_mm = tl.program_id(1) + tl.device_assert(tile_in_mm >= 0, "tile_in_mm < 0") + + tile_k, tile_n = _remap_xcd_tile_grid( + tile_in_mm, num_k_tiles, num_n_tiles, GROUP_SIZE=GROUP_SIZE + ) + + tl.device_assert(tile_k * BLOCK_SIZE_K >= 0, "tile_k * BLOCK_SIZE_K < 0") + tl.device_assert(tile_n * BLOCK_SIZE_N >= 0, "tile_n * BLOCK_SIZE_N < 0") + + offs_lhs_k = (tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)) % K + offs_rhs_n = (tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_m = tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + + if TRANS_LHS: + lhs_ptrs = lhs_ptr + offs_lhs_k[:, None] + (start_m + offs_m[None, :]) * K + else: + lhs_ptrs = lhs_ptr + offs_lhs_k[:, None] * M + (start_m + offs_m[None, :]) + + rhs_ptrs = rhs_ptr + (start_m + offs_m[:, None]) * N + offs_rhs_n[None, :] + + loop_m = tl.cdiv(m, BLOCK_SIZE_M) + m_divisible_by_block_m = m % BLOCK_SIZE_M == 0 + if not m_divisible_by_block_m: + loop_m -= 1 + + acc = tl.zeros((BLOCK_SIZE_K, BLOCK_SIZE_N), dtype=tl.float32) + # Initialize bias accumulator + bias_acc = tl.zeros((BLOCK_SIZE_K,), dtype=tl.float32) + + for _ in range(0, loop_m): + lhs = tl.load(lhs_ptrs) + rhs = tl.load(rhs_ptrs) + + acc = tl.dot(lhs, rhs, acc=acc) + + # Accumulate for bias gradient: sum lhs across M dimension + if COMPUTE_BIAS_GRAD and tile_n == 0: + bias_acc += tl.sum(lhs, axis=1) # [K, M] -> [K] + + if TRANS_LHS: + lhs_ptrs += BLOCK_SIZE_M * K + else: + lhs_ptrs += BLOCK_SIZE_M + + rhs_ptrs += BLOCK_SIZE_M * N + + if not m_divisible_by_block_m: + offs_lhs_k = ( + tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + ) % K + offs_rhs_n = ( + tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + ) % N + offs_m = loop_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + lhs = tl.load(lhs_ptrs, mask=offs_m[None, :] < m, other=0) + rhs = tl.load(rhs_ptrs, mask=offs_m[:, None] < m, other=0) + acc = tl.dot(lhs, rhs, acc=acc) + # Accumulate last chunk for bias gradient + if COMPUTE_BIAS_GRAD and tile_n == 0: + bias_acc += tl.sum(lhs, axis=1) + + acc = acc.to(out_ptr.type.element_ty) + + offs_out_k = tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + offs_out_n = tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + + out_ptrs = ( + out_ptr + g.to(tl.int64) * K * N + offs_out_k[:, None] * N + offs_out_n[None, :] + ) + + mask = (offs_out_k[:, None] < K) & (offs_out_n[None, :] < N) + if ACCUMULATE: + # Load existing values and add to them (like beta=1 in BLAS) + old_vals = tl.load(out_ptrs, mask=mask, other=0.0) + tl.store(out_ptrs, acc + old_vals, mask=mask) + else: + # Overwrite output (like beta=0 in BLAS) + tl.store(out_ptrs, acc, mask=mask) + + # Store bias gradient (only for first N tile, sum across all M) + if COMPUTE_BIAS_GRAD and tile_n == 0: + # Keep as float32 for atomic_add (bf16/fp16 not supported for atomics) + bias_grad_ptrs = bias_grad_ptr + g.to(tl.int64) * K + offs_out_k + # Use atomic add since multiple K-tiles may write to same expert's bias + tl.atomic_add(bias_grad_ptrs, bias_acc, mask=offs_out_k < K, sem="relaxed") diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/adapter.py b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..98c224244f27445384e0c2377d73516406927536 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/adapter.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Adapt AITER's Triton grouped GEMM to MegaBlocks' ``gmm`` calling convention. + +MegaBlocks (following tgale96/grouped_gemm) uses a single ``gmm`` entry point +with ``trans_a`` / ``trans_b`` flags: + +* ``trans_a=False, trans_b=False``: a(M,K) @ b(G,K,N) -> c(M,N) +* ``trans_a=False, trans_b=True`` : a(M,K) @ b(G,N,K)^T -> c(M,N) (dgrad) +* ``trans_a=True`` : a(M,K)^T @ b(M,N) per group -> c(G,K,N) (wgrad) + +AITER exposes these as two kernels: ``gmm`` ((M,K)@(G,K,N)->(M,N), transposition +of the 3D operand inferred from strides) and ``ptgmm`` ((K,M)@(M,N)->(G,K,N), +transposition of the 2D operand inferred from strides). +""" + +import torch + +from .gmm import gmm as _aiter_gmm +from .gmm import ptgmm as _aiter_ptgmm + + +def gmm(a, b, c, batch_sizes, trans_a=False, trans_b=False): + # AITER requires group sizes to be int32 and to live on the compute device. + group_sizes = batch_sizes.to(device=a.device, dtype=torch.int32) + + # AITER asserts exact strides: gmm wants lhs/rhs row-major (a transposed + # 3D operand must be exactly column-major), tgmm wants rhs row-major and + # lhs row/column-major. Make operands contiguous first so the transposed + # views have the precise strides the kernels expect. `.contiguous()` is a + # no-op when the tensor is already contiguous. + if trans_a: + # Weight gradient: a(M,K), b(M,N) -> c(G,K,N). + # Pass a transposed so AITER sees lhs(K,M) column-major (TRANS_LHS). + _aiter_ptgmm( + a.contiguous().transpose(0, 1), + b.contiguous(), + group_sizes, + preferred_element_type=c.dtype, + existing_out=c, + ) + else: + # trans_b contracts b's last dim: pass a column-major (G,K,N) view. + rhs = b.contiguous() + if trans_b: + rhs = rhs.transpose(1, 2) + _aiter_gmm( + a.contiguous(), + rhs, + group_sizes, + preferred_element_type=c.dtype, + existing_out=c, + ) + return c diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/configs.py b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/configs.py new file mode 100644 index 0000000000000000000000000000000000000000..9a4fe5617d8100869aa76dba9b7d22c7bcab814f --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/configs.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: MIT +# Tuned GMM configs vendored from ROCm/aiter (aiter/ops/triton/configs/). +# Inlined as a Python module so packaging always includes them. + +CONFIGS = {'gfx1250': {'gmm': {'default': {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_K': 64, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}}, 'ptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}}, 'nptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}}}, 'gfx942': {'gmm': {'default': {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_K': 64, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 304, 'num_warps': 8, 'num_stages': 1}}, 'ptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 304, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'GRID_DIM': 304, 'num_warps': 8, 'num_stages': 1}}, 'nptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}}}, 'gfx950': {'gmm': {'default': {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_K': 64, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}}, 'ptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}}, 'nptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}}}} diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/gmm.py b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/gmm.py new file mode 100644 index 0000000000000000000000000000000000000000..e30c9326c6d4e4836d1303e2761ea2440a7f4750 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/gmm.py @@ -0,0 +1,567 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved. + + +# Imports. +# ------------------------------------------------------------------------------ + +# PyTorch +import torch +from torch import Tensor + +# Triton +import triton + +# AITER: GMM utility functions +from .utils.gmm_common import ( + DTYPE, + is_power_of_2, + check_input_device_dtype, + check_bias_shape_stride, + get_gmm_shape, + get_gmm_output, + get_gmm_transposition, + get_tgmm_shape, + get_tgmm_output, + get_tgmm_bias_grad, + get_tgmm_transposition, +) + +# AITER: GMM Triton kernels +from ._triton_kernels.gmm import ( + gmm_kernel, + tgmm_persistent_kernel, + tgmm_non_persistent_kernel, + get_config, +) + +# GMM PyTorch wrapper. +# ------------------------------------------------------------------------------ + + +def _gmm_grid( + N: int, + block_size_m: int, + block_size_n: int, + group_sizes: Tensor, + grid_dim: int, +) -> tuple[int]: + assert N > 0, f"N must be positive, it's {N}." + assert is_power_of_2( + block_size_m + ), f"M-dimension tile size must be a power of 2 (it's {block_size_m})." + assert is_power_of_2( + block_size_n + ), f"N-dimension tile size must be a power of 2 (it's {block_size_n})." + assert torch.all(group_sizes >= 0).item(), "All group_sizes must be non-negative." + assert grid_dim > 0, f"Grid dimension must be positive (it's {grid_dim})." + num_m_tiles = (group_sizes + block_size_m - 1) // block_size_m + assert torch.all(num_m_tiles >= 0).item(), "All num_m_tiles must be non-negative." + num_n_tiles = triton.cdiv(N, block_size_n) + assert num_n_tiles > 0, f"num_n_tiles must be positive, it's {num_n_tiles}." + num_tiles = torch.sum(num_m_tiles * num_n_tiles).item() + assert num_tiles > 0, f"num_tiles must be positive, it's {num_tiles}." + num_programs = int(min(grid_dim, num_tiles)) + assert num_programs > 0, f"num_programs must be positive, it's {num_programs}." + return (num_programs,) + + +def gmm( + lhs: Tensor, + rhs: Tensor, + group_sizes: Tensor, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, + config: dict[str, int] | None = None, + bias: Tensor | None = None, +) -> Tensor: + """ + Perform Group Matrix Multiplication (GMM): out = lhs @ rhs + bias + + lhs rows are divided into G groups. Each group of lhs rows is matrix multiplied with a plane of + rhs 3D tensor and then stored in a slice of out. In PyTorch parlance, it can be implemented as + follows for a given group g: + out[group_start:group_end, :] = lhs[group_start:group_end, :] @ rhs[g] + bias[g] + + The size of each group, and their respective start and end positions are specified by + group_sizes tensor. For instance, suppose that group_sizes = [3, 2, 4, 1]. In this particular + case we have 4 groups. The 1st group starts at 0 and ends at 2, the second group starts at 3 and + ends at 4, the third group starts at 5 and ends at 8, and the fourth and final group consists of + just the 10th (last) row of lhs. + + Parameters + ---------- + lhs : torch.Tensor + Left-hand side 2D input tensor. Shape: (M, K). + lhs data type must be torch.float16 or torch.bfloat16, and must match rhs data type. + lhs must be on the same device of rhs and group_sizes. + rhs : torch.Tensor + Right-hand side 3D input tensor. Shape: (G, K, N). + rhs data type must be torch.float16 or torch.bfloat16, and must match lhs data type. + rhs must be on the same device of lhs and group_sizes. + group_sizes : torch.Tensor + 1D input tensor describing group sizes. Shape: (G,). + group_sizes data type must be torch.int32 and all its elements must be non-negative. + group_sizes must be on the same device of lhs and rhs. + preferred_element_type : torch.dtype, optional + Desired data type for output tensor. Default is torch.bfloat16. + Supported output types are torch.float16 and torch.bfloat16. + existing_out : torch.Tensor or None, optional + Preallocated output tensor. Default is None. + If provided, results are written into this tensor. Otherwise, a new output tensor is + allocated. + If provided then it must have shape (M, N), its data type must match preferred_element_type + and it must be on the same device of other input tensors. + config : dict[str, int] or None, optional + Optional dictionary with kernel metaparameters. If absent, config will be queried from + internal tuning database. + bias : torch.Tensor or None, optional + Optional bias tensor. Shape: (G, N). + If provided, bias data type must match lhs and rhs data type, and bias must be on the same + device as other input tensors. Each group g adds bias[g] to the output. + + Returns + ------- + torch.Tensor + The computed output 2D tensor. Shape: (M, N). + Output tensor data type is given by preferred_element_type. + If existing_out is provided then existing_out is also returned. + + Implementation Notes + -------------------- + - GMM is implemented with a persistent Triton kernel. + - lhs must be row-major (lhs.stride() == (K, 1)). + - rhs can be row-major (rhs.stride() == (K * N, N, 1)) or column-major (rhs.stride() == + (K * N, 1, K)). If rhs is row-major then kernel parameter TRANS_RHS == False, this is useful + for implementing forward pass. If rhs is column-major then kernel parameter TRANS_RHS == True, + this is useful for computing the lhs derivative in the backward pass, while fusing the + transposition. + - out must be row-major (out.stride() == (N, 1)). + - bias must be row-major (bias.stride() == (N, 1)) if provided. + """ + use_bias = bias is not None + check_input_device_dtype(lhs, rhs, group_sizes, bias) + + M, K, N, G = get_gmm_shape(lhs, rhs, group_sizes) + + if use_bias: + check_bias_shape_stride(bias, G, N) + + out = get_gmm_output( + M, + N, + device=lhs.device, + preferred_element_type=preferred_element_type, + existing_out=existing_out, + ) + + trans_rhs, _ = get_gmm_transposition(lhs, rhs, out) + + if config is None: + config = get_config("gmm", M, K, N, G) + + assert all( + key in config + and isinstance(config[key], int) + and ( + is_power_of_2(config[key]) + if key.startswith("BLOCK_SIZE_") + else config[key] > 0 + ) + for key in { + "BLOCK_SIZE_M", + "BLOCK_SIZE_K", + "BLOCK_SIZE_N", + "GROUP_SIZE", + "GRID_DIM", + } + ), "Invalid GMM kernel config." + + grid = _gmm_grid( + N, + config["BLOCK_SIZE_M"], + config["BLOCK_SIZE_N"], + group_sizes, + config["GRID_DIM"], + ) + + # fmt: off + gmm_kernel[grid]( + # Tensor pointers: + lhs, rhs, group_sizes, out, bias, + # Tensor shapes: + M, K, N, G, + # Meta-parameters: + TRANS_RHS=trans_rhs, + USE_BIAS=use_bias, + **config, + ) + # fmt: on + + return out + + +# Persistent TGMM PyTorch wrapper. +# ------------------------------------------------------------------------------ + + +def _ptgmm_grid( + K: int, + N: int, + G: int, + block_size_k: int, + block_size_n: int, + grid_dim: int, +) -> tuple[int]: + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}." + assert G > 0, f"G must be positive, it's {G}." + assert is_power_of_2( + block_size_k + ), f"K-dimension tile size must be a power of 2 (it's {block_size_k})." + assert is_power_of_2( + block_size_n + ), f"N-dimension tile size must be a power of 2 (it's {block_size_n})." + assert grid_dim > 0, f"Grid dimension must be positive (it's {grid_dim})." + num_k_tiles = triton.cdiv(K, block_size_k) + assert num_k_tiles > 0, f"num_k_tiles must be positive, it's {num_k_tiles}." + num_n_tiles = triton.cdiv(N, block_size_n) + assert num_n_tiles > 0, f"num_n_tiles must be positive, it's {num_n_tiles}." + num_tiles = G * num_k_tiles * num_n_tiles + assert num_tiles > 0, f"num_tiles must be positive, it's {num_tiles}." + num_programs = min(grid_dim, num_tiles) + assert num_programs > 0, f"num_programs must be positive, it's {num_programs}." + return (num_programs,) + + +def ptgmm( + lhs: Tensor, + rhs: Tensor, + group_sizes: Tensor, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, + config: dict[str, int] | None = None, + bias_grad: Tensor | None = None, + accumulate: bool = False, +) -> Tensor: + """ + Perform a Group Matrix Multiplication (GMM) variant: out = lhs @ rhs + + lhs columns and rhs rows are divided into G groups. Each group of lhs is matrix multiplied with + the respective group of rhs and then stored in a plane of the output 3D tensor. In PyTorch + parlance, it can be implemented as follows for a given group g: + out[g] = lhs[:, group_start:group_end] @ rhs[group_start:group_end, :] + + The 't' in the operator name derives from MaxText implementation + (https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/kernels/megablox/gmm.py), + which served as the initial inspiration for this one. TGMM differs from GMM in terms of tensor + shapes. GMM does (M, K) @ (G, K, N) = (M, N) while TGMM does (K, M) @ (M, N) = (G, K, N). + + The 'p' in the operator name means that it is implemented with a persistent kernel. There is + also the non-persistent variation, which is implemented with a regular kernel. Please take a + look at nptgmm operator. Both ptgmm and nptgmm implement the same computation, choosing one or + the other is a matter of performance for the target workload. + + Parameters + ---------- + lhs : torch.Tensor + Left-hand side 2D input tensor. Shape: (K, M). + lhs data type must be torch.float16 or torch.bfloat16, and must match rhs data type. + lhs must be on the same device of rhs and group_sizes. + rhs : torch.Tensor + Right-hand side 2D input tensor. Shape: (M, N). + rhs data type must be torch.float16 or torch.bfloat16, and must match lhs data type. + rhs must be on the same device of lhs and group_sizes. + group_sizes : torch.Tensor + 1D input tensor describing group sizes. Shape: (G,). + group_sizes data type must be torch.int32 and all its elements must be non-negative. + group_sizes must be on the same device of lhs and rhs. + preferred_element_type : torch.dtype, optional + Desired data type for output tensor. Default is torch.bfloat16. + Supported output types are torch.float16 and torch.bfloat16. + existing_out : torch.Tensor or None, optional + Preallocated output tensor. Default is None. + If provided, results are written into this tensor. Otherwise, a new output tensor is + allocated. + If provided then it must have shape (G, K, N), its data type must match + preferred_element_type and it must be on the same device of other input tensors. + config : dict[str, int] or None, optional + Optional dictionary with kernel metaparameters. If absent, config will be queried from + internal tuning database. + bias_grad : torch.Tensor or None, optional + Optional bias gradient output tensor. Shape: (G, K). + If provided, the kernel will compute the bias gradient and write it to this tensor. + bias_grad must be torch.float32 (kernel uses atomic_add which requires float32), + accumulate : bool, optional + Whether to accumulate into existing output tensor values. Default is False. + If False, output will be overwritten with fresh computation. + If True, results will be added to existing output tensor values. + + Returns + ------- + torch.Tensor + The computed output 3D tensor. Shape: (G, K, N). + Output tensor data type is given by preferred_element_type. + If existing_out is provided then existing_out is also returned. + + Implementation Notes + -------------------- + - PTGMM is implemented with a persistent Triton kernel. + - lhs can be row-major (lhs.stride() == (M, 1)) or column-major (lhs.stride() == (1, K)). If lhs + is row-major then kernel parameter TRANS_LHS == False. If lhs is column-major then kernel + parameter TRANS_LHS == True, this is useful for computing the rhs derivative in the backward + pass, while fusing the transposition. + - rhs must be row-major (rhs.stride() == (N, 1)). + - out must be row-major (out.stride() == (K * N, N, 1)). + """ + check_input_device_dtype(lhs, rhs, group_sizes) + + M, K, N, G = get_tgmm_shape(lhs, rhs, group_sizes) + + out = get_tgmm_output( + K, + N, + G, + device=lhs.device, + preferred_element_type=preferred_element_type, + existing_out=existing_out, + ) + + trans_lhs, _ = get_tgmm_transposition(lhs, rhs, out) + + if config is None: + config = get_config("ptgmm", M, K, N, G, accumulate) + + assert all( + key in config + and isinstance(config[key], int) + and ( + is_power_of_2(config[key]) + if key.startswith("BLOCK_SIZE_") + else config[key] > 0 + ) + for key in { + "BLOCK_SIZE_M", + "BLOCK_SIZE_K", + "BLOCK_SIZE_N", + "GROUP_SIZE", + "GRID_DIM", + } + ), "Invalid PTGMM kernel config." + + # Bias gradient handling. + # ----------------------- + # Get or validate bias gradient tensor. + compute_bias_grad = bias_grad is not None + bias_grad_ptr = get_tgmm_bias_grad( + K, + G, + device=lhs.device, + existing_bias_grad=bias_grad, + ) + + grid = _ptgmm_grid( + K, + N, + G, + config["BLOCK_SIZE_K"], + config["BLOCK_SIZE_N"], + config["GRID_DIM"], + ) + + # fmt: off + tgmm_persistent_kernel[grid]( + # Tensor pointers: + lhs, rhs, group_sizes, out, bias_grad_ptr, + # Tensor shapes: + M, K, N, G, + # Meta-parameters: + TRANS_LHS=trans_lhs, + COMPUTE_BIAS_GRAD=compute_bias_grad, + ACCUMULATE=accumulate, + **config, + ) + # fmt: on + + return out + + +# Regular non-persistent TGMM PyTorch wrapper. +# ------------------------------------------------------------------------------ + + +def _nptgmm_grid( + K: int, + N: int, + G: int, + block_size_k: int, + block_size_n: int, +) -> tuple[int, int]: + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}." + assert G > 0, f"G must be positive, it's {G}." + assert is_power_of_2( + block_size_k + ), f"K-dimension tile size must be a power of 2 (it's {block_size_k})." + assert is_power_of_2( + block_size_n + ), f"N-dimension tile size must be a power of 2 (it's {block_size_n})." + num_k_tiles = triton.cdiv(K, block_size_k) + assert num_k_tiles > 0, f"num_k_tiles must be positive, it's {num_k_tiles}." + num_n_tiles = triton.cdiv(N, block_size_n) + assert num_n_tiles > 0, f"num_n_tiles must be positive, it's {num_n_tiles}." + num_tiles_per_mm = num_k_tiles * num_n_tiles + assert ( + num_tiles_per_mm > 0 + ), f"num_tiles_per_mm must be positive, it's {num_tiles_per_mm}." + return (G, num_tiles_per_mm) + + +def nptgmm( + lhs: Tensor, + rhs: Tensor, + group_sizes: Tensor, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, + config: dict[str, int] | None = None, + bias_grad: Tensor | None = None, + accumulate: bool = False, +) -> Tensor: + """ + Perform a Group Matrix Multiplication (GMM) variant: out = lhs @ rhs + + lhs columns and rhs rows are divided into G groups. Each group of lhs is matrix multiplied with + the respective group of rhs and then stored in a plane of the output 3D tensor. In PyTorch + parlance, it can be implemented as follows for a given group g: + out[g] = lhs[:, group_start:group_end] @ rhs[group_start:group_end, :] + + The 't' in the operator name derives from MaxText implementation + (https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/kernels/megablox/gmm.py), + which served as the initial inspiration for this one. TGMM differs from GMM in terms of tensor + shapes. GMM does (M, K) @ (G, K, N) = (M, N) while TGMM does (K, M) @ (M, N) = (G, K, N). + + The 'np' in the operator name means that it is implemented with a non-persistent, i.e. regular + kernel. There is also the persistent variation, which is implemented with a persistent kernel. + Please take a look at ptgmm operator. Both nptgmm and ptgmm implement the same computation, + choosing one or the other is a matter of performance for the target workload. + + Parameters + ---------- + lhs : torch.Tensor + Left-hand side 2D input tensor. Shape: (K, M). + lhs data type must be torch.float16 or torch.bfloat16, and must match rhs data type. + lhs must be on the same device of rhs and group_sizes. + rhs : torch.Tensor + Right-hand side 2D input tensor. Shape: (M, N). + rhs data type must be torch.float16 or torch.bfloat16, and must match lhs data type. + rhs must be on the same device of lhs and group_sizes. + group_sizes : torch.Tensor + 1D input tensor describing group sizes. Shape: (G,). + group_sizes data type must be torch.int32 and all its elements must be non-negative. + group_sizes must be on the same device of lhs and rhs. + preferred_element_type : torch.dtype, optional + Desired data type for output tensor. Default is torch.bfloat16. + Supported output types are torch.float16 and torch.bfloat16. + existing_out : torch.Tensor or None, optional + Preallocated output tensor. Default is None. + If provided, results are written into this tensor. Otherwise, a new output tensor is + allocated. + If provided then it must have shape (G, K, N), its data type must match + preferred_element_type and it must be on the same device of other input tensors. + config : dict[str, int] or None, optional + Optional dictionary with kernel metaparameters. If absent, config will be queried from + internal tuning database. + bias_grad : torch.Tensor or None, optional + Optional bias gradient output tensor. Shape: (G, K). + If provided, the kernel will compute the bias gradient and write it to this tensor. + bias_grad must be torch.float32 (kernel uses atomic_add which requires float32), + accumulate : bool, optional + Whether to accumulate into existing output tensor values. Default is False. + If False, output will be overwritten with fresh computation. + If True, results will be added to existing output tensor values. + + Returns + ------- + torch.Tensor + The computed output 3D tensor. Shape: (G, K, N). + Output tensor data type is given by preferred_element_type. + If existing_out is provided then existing_out is also returned. + + Implementation Notes + -------------------- + - NPTGMM is implemented with a non-persistent regular Triton kernel. + - lhs can be row-major (lhs.stride() == (M, 1)) or column-major (lhs.stride() == (1, K)). If lhs + is row-major then kernel parameter TRANS_LHS == False. If lhs is column-major then kernel + parameter TRANS_LHS == True, this is useful for computing the rhs derivative in the backward + pass, while fusing the transposition. + - rhs must be row-major (rhs.stride() == (N, 1)). + - out must be row-major (out.stride() == (K * N, N, 1)). + """ + check_input_device_dtype(lhs, rhs, group_sizes) + + M, K, N, G = get_tgmm_shape(lhs, rhs, group_sizes) + + out = get_tgmm_output( + K, + N, + G, + device=lhs.device, + preferred_element_type=preferred_element_type, + existing_out=existing_out, + ) + + trans_lhs, _ = get_tgmm_transposition(lhs, rhs, out) + + # Bias gradient handling. + # ----------------------- + # Get or validate bias gradient tensor. + compute_bias_grad = bias_grad is not None + bias_grad_ptr = get_tgmm_bias_grad( + K, + G, + device=lhs.device, + existing_bias_grad=bias_grad, + ) + + if config is None: + config = get_config("nptgmm", M, K, N, G, accumulate) + + assert all( + key in config + and isinstance(config[key], int) + and ( + is_power_of_2(config[key]) + if key.startswith("BLOCK_SIZE_") + else config[key] > 0 + ) + for key in { + "BLOCK_SIZE_M", + "BLOCK_SIZE_K", + "BLOCK_SIZE_N", + "GROUP_SIZE", + } + ), "Invalid NPTGMM kernel config." + + grid = _nptgmm_grid( + K, + N, + G, + config["BLOCK_SIZE_K"], + config["BLOCK_SIZE_N"], + ) + + # fmt: off + tgmm_non_persistent_kernel[grid]( + # Tensor pointers: + lhs, rhs, group_sizes, out, bias_grad_ptr, + # Tensor shapes: + M, K, N, G, + # Meta-parameters: + TRANS_LHS=trans_lhs, + COMPUTE_BIAS_GRAD=compute_bias_grad, + ACCUMULATE=accumulate, + **config, + ) + # fmt: on + + return out diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/_triton/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/_triton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/_triton/arch_info.py b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/_triton/arch_info.py new file mode 100644 index 0000000000000000000000000000000000000000..3f6c88581a64044518125623f116082c53bd5474 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/_triton/arch_info.py @@ -0,0 +1,46 @@ +import triton + +# Detect the GPU arch lazily: querying the triton driver at import time fails +# in headless environments (e.g. the kernel-builder ABI check sandbox has no +# GPU), and the original JAX fallback pulled in an unrelated runtime dep. The +# arch is only actually needed when a GMM kernel is dispatched, so resolve and +# cache on first call. +_CACHED_ARCH = None + + +def get_arch(): + global _CACHED_ARCH + if _CACHED_ARCH is not None: + return _CACHED_ARCH + try: + _CACHED_ARCH = triton.runtime.driver.active.get_current_target().arch + except RuntimeError: + try: + from jax._src.lib import gpu_triton as triton_kernel_call_lib + _CACHED_ARCH = triton_kernel_call_lib.get_arch_details("0").split(":")[0] + except ImportError as e: + raise RuntimeError( + "Cannot determine GPU arch: triton driver is inactive and " + "JAX is not available. A GPU is required for grouped GEMM." + ) from e + return _CACHED_ARCH + + +def is_gluon_avail(): + return get_arch() in ("gfx950", "gfx1250") + + +def is_fp4_avail(): + return get_arch() in ("gfx950", "gfx1250") + + +def is_fp8_avail(): + return get_arch() in ("gfx942", "gfx950", "gfx1250", "gfx1200", "gfx1201") + + +def is_mx_scale_preshuffling_avail(): + return get_arch() in ("gfx950", "gfx1250") + + +def is_tdm_avail(): + return get_arch() in ("gfx1250",) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/_triton/pid_preprocessing.py b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/_triton/pid_preprocessing.py new file mode 100644 index 0000000000000000000000000000000000000000..99792bb3ba2fab8fff223bba733ced1eb6e6df53 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/_triton/pid_preprocessing.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: MIT + +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import triton +import triton.language as tl + + +@triton.jit +def remap_xcd_chunked( + pid, GRID_MN, NUM_XCDS: tl.constexpr = 8, CHUNK_SIZE: tl.constexpr = 2 +): + # Compute current XCD and local PID + xcd = pid % NUM_XCDS + # distribute the modulo pids in round robin + if pid > (GRID_MN // (NUM_XCDS * CHUNK_SIZE)) * (NUM_XCDS * CHUNK_SIZE): + return pid + local_pid = pid // NUM_XCDS + # Calculate chunk index and position within chunk + chunk_idx = local_pid // CHUNK_SIZE + pos_in_chunk = local_pid % CHUNK_SIZE + # Calculate new PID + new_pid = chunk_idx * NUM_XCDS * CHUNK_SIZE + xcd * CHUNK_SIZE + pos_in_chunk + return new_pid + + +@triton.jit +def remap_xcd(pid, GRID_MN, NUM_XCDS: tl.constexpr = 8): + ## pid remapping on xcds + # Number of pids per XCD in the new arrangement + pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS + # When GRID_MN cannot divide NUM_XCDS, some xcds will have + # pids_per_xcd pids, the other will have pids_per_xcd - 1 pids. + # We calculate the number of xcds that have pids_per_xcd pids as + # tall_xcds + tall_xcds = GRID_MN % NUM_XCDS + tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds + # Compute current XCD and local pid within the XCD + xcd = pid % NUM_XCDS + local_pid = pid // NUM_XCDS + # Calculate new pid based on the new grouping + # Note that we need to consider the following two cases: + # 1. the current pid is on a tall xcd + # 2. the current pid is on a short xcd + if xcd < tall_xcds: + pid = xcd * pids_per_xcd + local_pid + else: + pid = ( + tall_xcds * pids_per_xcd + + (xcd - tall_xcds) * (pids_per_xcd - 1) + + local_pid + ) + + return pid + + +@triton.jit +def pid_grid(pid: int, num_pid_m: int, num_pid_n: int, GROUP_SIZE_M: tl.constexpr = 1): + """ + Maps 1D pid to 2D grid coords (pid_m, pid_n). + + Args: + - pid: 1D pid + - num_pid_m: grid m size + - num_pid_n: grid n size + - GROUP_SIZE_M: tl.constexpr: default is 1 + """ + if GROUP_SIZE_M == 1: + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + else: + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + tl.assume(group_size_m >= 0) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + return pid_m, pid_n + + +@triton.jit +def pid_grid_3d(pid: int, num_pid_m: int, num_pid_n: int, num_pid_k): + """ + Maps 1D pid to 3D grid coords (pid_m, pid_n, pid_k). + Args: + - pid: 1D pid + - num_pid_m: grid m size + - num_pid_n: grid n size + - num_pid_k: grid k size + + Returns: + - pid_m, pid_n, pid_k: 3D grid coordinates + """ + pid_m = pid % num_pid_m + pid_n = (pid // num_pid_m) % num_pid_n + pid_k = pid // (num_pid_m * num_pid_n) % num_pid_k + + return pid_m, pid_n, pid_k diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/gmm_common.py b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/gmm_common.py new file mode 100644 index 0000000000000000000000000000000000000000..153dee65b50ab5f833262481889d2184d1ca639f --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/gmm_common.py @@ -0,0 +1,752 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved. + +# Imports. +# ------------------------------------------------------------------------------ + +# PyTorch +import torch +from torch import Tensor + +# AITER: logging +from .logger import AiterTritonLogger + +_LOGGER: AiterTritonLogger = AiterTritonLogger() + + +# Supported data types. +# ------------------------------------------------------------------------------ + +# Supported data types, as strings. +SUPPORTED_DTYPES_STR: set[str] = {"fp16", "bf16"} + + +# Convert string data type to PyTorch data type. +def dtype_from_str(dtype_str: str) -> torch.dtype: + dtype_str = dtype_str.strip().lower() + dtype_str = dtype_str[1:] if dtype_str[0] in {"i", "o"} else dtype_str + assert ( + dtype_str in SUPPORTED_DTYPES_STR + ), "String data type isn't in set of supported string data types." + return {"fp16": torch.float16, "bf16": torch.bfloat16}[dtype_str] + + +# Supported data types, as PyTorch types. +SUPPORTED_DTYPES: set[torch.dtype] = { + dtype_from_str(dtype_str) for dtype_str in SUPPORTED_DTYPES_STR +} + + +# Convert PyTorch data type to string data type. +def str_from_dtype(dtype: torch.dtype) -> str: + assert ( + dtype in SUPPORTED_DTYPES + ), "PyTorch data type isn't in set of supported PyTorch data types." + return {torch.float16: "fp16", torch.bfloat16: "bf16"}[dtype] + + +# Default data type, as string. +DTYPE_STR: str = "bf16" +assert ( + DTYPE_STR in SUPPORTED_DTYPES_STR +), "Default string data type isn't in set of supported string data types." + + +# Default data type, as PyTorch type. +DTYPE: torch.dtype = dtype_from_str(DTYPE_STR) + + +# Other defaults. +# ------------------------------------------------------------------------------ + +# Default device. +DEVICE: torch.device | str = "cuda" + +# Default RNG seed for input generation. +RNG_SEED: int = 0 + +# Default number of group sizes. +NUM_GROUP_SIZES: int = 1 + +# Default transposition (NN). +TRANS_LHS: bool = False +TRANS_RHS: bool = False + + +# Parameter checking functions. +# ------------------------------------------------------------------------------ + + +def is_power_of_2(x: int) -> bool: + return (x > 0) and (x & (x - 1) == 0) + + +def check_input_device_dtype( + lhs: Tensor, rhs: Tensor, group_sizes: Tensor, bias: Tensor | None = None +) -> None: + assert ( + lhs.device == rhs.device == group_sizes.device + ), f"All input tensors must be in the same device (lhs = {lhs.device}, rhs = {rhs.device}, group_sizes = {group_sizes.device})." + assert ( + lhs.dtype == rhs.dtype + ), f"lhs and rhs types must match (lhs = {lhs.dtype}, rhs = {rhs.dtype})." + assert group_sizes.dtype == torch.int32, "group_sizes type must be int32." + + if bias is not None: + assert ( + bias.device == lhs.device + ), f"bias must be on the same device as lhs (bias = {bias.device}, lhs = {lhs.device})." + assert ( + bias.dtype == lhs.dtype + ), f"bias dtype must match lhs dtype (bias = {bias.dtype}, lhs = {lhs.dtype})." + + +def check_bias_shape_stride(bias: Tensor, G: int, N: int) -> None: + assert bias.shape == ( + G, + N, + ), f"bias must have shape (G, N) = ({G}, {N}), got {bias.shape}." + assert bias.stride() == (N, 1), "bias must be row-major (bias.stride() == (N, 1))." + + +# Generation of group sizes. +# ------------------------------------------------------------------------------ + + +# Probabilities for generating random group sizes. +UNUSED_TOKENS_PROB: float = 0.0 +UNUSED_EXPERTS_PROB: float = 0.1 + + +def gen_uniform_group_sizes( + M: int, + G: int, + device: torch.device | str = DEVICE, +) -> Tensor: + assert M >= 0, f"Number of tokens M must be non-negative (it's {M})." + assert G > 0, f"Number of experts G must be positive (it's {G})." + + base = M // G + remainder = M % G + group_sizes = torch.full((G,), base, dtype=torch.int32, device=device) + if remainder > 0: + group_sizes[:remainder] += 1 + + assert ( + len(group_sizes) == G + ), f"Group sizes don't have {G} elements (it's {len(group_sizes)})." + assert torch.all(group_sizes >= 0).item(), "All group sizes must be non-negative." + assert ( + torch.sum(group_sizes).item() == M + ), f"Group sizes don't add up to total tokens {M}." + assert group_sizes.dtype == torch.int32, "Group sizes must be int32." + + return group_sizes + + +def gen_group_sizes( + M: int, + G: int, + device: torch.device | str = DEVICE, + rng_seed: int | None = RNG_SEED, + unused_tokens_prob: float = UNUSED_TOKENS_PROB, + unused_experts_prob: float = UNUSED_EXPERTS_PROB, +) -> Tensor: + assert M >= 0, f"Number of tokens M must be non-negative (it's {M})." + assert G > 0, f"Number of experts G must be positive (it's {G})." + assert ( + 0 <= unused_tokens_prob <= 1 + ), f"Probability of unused tokens must be in [0, 1] interval (it's {unused_tokens_prob})." + assert ( + 0 <= unused_experts_prob <= 1 + ), f"Probability of unused experts must be in [0, 1] interval (it's {unused_experts_prob})." + + if rng_seed is not None: + torch.manual_seed(rng_seed) + + if unused_tokens_prob > 0: + # Optionally drop tokens to simulate routing sparsity, some tokens may not be routed. + num_unused_tokens = M + while num_unused_tokens == M: + num_unused_tokens = int( + torch.binomial( + torch.tensor(float(M), device=device), + torch.tensor(unused_tokens_prob, device=device), + ).item() + ) + else: + num_unused_tokens = 0 + num_used_tokens = M - num_unused_tokens + assert ( + num_unused_tokens >= 0 + ), f"Number of unused tokens must be non-negative (it's {num_unused_tokens})." + assert ( + num_used_tokens > 0 + ), f"Number of used tokens must be positive (it's {num_used_tokens})." + assert ( + num_used_tokens + num_unused_tokens == M + ), f"Unused + used tokens don't add up total tokens ({num_used_tokens} + {num_unused_tokens} != {M})." + + if num_unused_tokens > 0: + _LOGGER.debug( + f"Group sizes generation: dropped {num_unused_tokens} token{'s' if num_unused_tokens > 1 else ''}.", + ) + + if unused_experts_prob > 0: + # Some experts may have zero tokens assigned to them. + num_used_experts = 0 + while num_used_experts == 0: + used_experts = torch.nonzero( + torch.rand((G,), device=device) >= unused_experts_prob + ).squeeze() + num_used_experts = used_experts.numel() + else: + used_experts = torch.arange(0, G, device=device) + num_used_experts = G + num_unused_experts = G - num_used_experts + assert ( + num_unused_experts >= 0 + ), f"Number of unused experts must be non-negative (it's {num_unused_experts})." + assert ( + num_used_experts >= 1 + ), f"At least one expert must be used (it's {num_used_experts})." + assert ( + num_unused_experts + num_used_experts == G + ), f"Unused + used experts don't add up total experts ({num_unused_experts} + {num_used_experts} != {G})." + + if num_unused_experts > 0: + _LOGGER.debug( + f"Group sizes generation: dropped {num_unused_experts} expert{'s' if num_unused_experts > 1 else ''}.", + ) + + group_sizes = torch.bincount( + used_experts[ + torch.randint(low=0, high=num_used_experts, size=(num_used_tokens,)) + ], + minlength=G, + ).to(torch.int32) + + assert ( + len(group_sizes) == G + ), f"Group sizes don't have {G} elements (it's {len(group_sizes)})." + assert torch.all(group_sizes >= 0).item(), "All group sizes must be non-negative." + assert ( + torch.sum(group_sizes).item() == num_used_tokens + ), f"Group sizes don't add up to used tokens {num_used_tokens}." + assert group_sizes.dtype == torch.int32, "Group sizes must be int32." + + return group_sizes + + +def gen_multiple_group_sizes( + num_group_sizes: int, + M: int, + G: int, + device: torch.device | str = DEVICE, + rng_seed: int | None = RNG_SEED, + unused_tokens_prob: float = UNUSED_TOKENS_PROB, + unused_experts_prob: float = UNUSED_EXPERTS_PROB, + group_sizes_0: Tensor | None = None, +) -> list[Tensor]: + assert ( + num_group_sizes > 0 + ), f"Number of group sizes to be generated must be positive, it's {num_group_sizes}." + multiple_group_sizes = [ + gen_group_sizes( + M, + G, + device=device, + rng_seed=rng_seed if g == 0 else None, + unused_tokens_prob=unused_tokens_prob, + unused_experts_prob=unused_experts_prob, + ) + for g in range( + num_group_sizes if group_sizes_0 is None else num_group_sizes - 1 + ) + ] + if group_sizes_0 is not None: + multiple_group_sizes.insert(0, group_sizes_0) + assert ( + len(multiple_group_sizes) == num_group_sizes + ), f"Expecting {num_group_sizes} distinct group sizes (it's {len(multiple_group_sizes)})." + return multiple_group_sizes + + +# GMM helpers: tensor generation. +# ------------------------------------------------------------------------------ + + +def gen_gmm_input( + M: int, + K: int, + N: int, + G: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, + trans_rhs: bool = TRANS_RHS, + rng_seed: int | None = RNG_SEED, + unif_group_sizes: bool = False, +) -> tuple[Tensor, Tensor, Tensor]: + assert M > 0, f"Number of lhs rows M must be positive (M = {M})." + assert K > 0, f"Number of lhs columns / rhs rows K must be positive (K = {K})." + assert N > 0, f"Number of rhs columns N must be positive (N = {N})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + if rng_seed is not None: + torch.manual_seed(rng_seed) + + lhs = torch.randn((M, K), dtype=torch.float32, device=device) + lhs = lhs.to(preferred_element_type) + + if trans_rhs: + rhs = torch.randn((G, N, K), dtype=torch.float32, device=device).permute( + 0, 2, 1 + ) + else: + rhs = torch.randn((G, K, N), dtype=torch.float32, device=device) + rhs = rhs.to(preferred_element_type) + + group_sizes = ( + gen_uniform_group_sizes(M, G, device=device) + if unif_group_sizes + else gen_group_sizes(M, G, device=device, rng_seed=None) + ) + + return lhs, rhs, group_sizes + + +def gen_gmm_output( + M: int, + N: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, +) -> Tensor: + assert M > 0, f"Number of out rows M must be positive (M = {M})." + assert N > 0, f"Number of out columns N must be positive (N = {N})." + + out = torch.empty((M, N), dtype=preferred_element_type, device=device) + + return out + + +def gen_gmm_tensors( + M: int, + K: int, + N: int, + G: int, + num_group_sizes: int, + device: torch.device | str = DEVICE, + input_type: torch.dtype = DTYPE, + output_type: torch.dtype = DTYPE, + trans_lhs: bool = False, + trans_rhs: bool = TRANS_RHS, + rng_seed: int | None = RNG_SEED, + unif_group_sizes: bool = False, + use_bias: bool = False, +) -> tuple[Tensor, Tensor, list[Tensor], Tensor, Tensor | None]: + lhs, rhs, group_sizes_0 = gen_gmm_input( + M, + K, + N, + G, + device=device, + preferred_element_type=input_type, + trans_rhs=trans_rhs, + rng_seed=rng_seed, + unif_group_sizes=unif_group_sizes, + ) + multiple_group_sizes = gen_multiple_group_sizes( + num_group_sizes, M, G, device=device, rng_seed=None, group_sizes_0=group_sizes_0 + ) + out = gen_gmm_output(M, N, device=device, preferred_element_type=output_type) + bias = None + if use_bias: + torch.manual_seed(rng_seed + 1000) # Different seed for bias + bias = torch.randn(G, N, dtype=input_type, device=device) + + return lhs, rhs, multiple_group_sizes, out, bias + + +# GMM helpers: get information from tensors. +# ------------------------------------------------------------------------------ + + +def get_gmm_shape( + lhs: Tensor, rhs: Tensor, group_sizes: Tensor +) -> tuple[int, int, int, int]: + assert lhs.dim() == 2, f"lhs must have 2 dimensions (it's {lhs.dim()})." + assert rhs.dim() == 3, f"rhs must have 3 dimensions (it's {rhs.dim()})." + assert ( + group_sizes.dim() == 1 + ), f"group_sizes must have 1 dimension (it's {group_sizes.dim()})." + + M, lhs_k = lhs.shape + rhs_g, rhs_k, N = rhs.shape + group_sizes_g = group_sizes.shape[0] + + assert ( + lhs_k == rhs_k + ), f"K dimension of lhs and rhs don't match (lhs = {lhs_k}, rhs = {rhs_k})." + K = lhs_k + assert ( + rhs_g == group_sizes_g + ), f"G dimension of rhs and group_sizes don't match (rhs = {rhs_g}, group_sizes = {group_sizes_g})." + G = rhs_g + + assert M > 0, f"M must be positive, it's {M}." + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}" + assert G > 0, f"G must be positive, it's {G}" + + return M, K, N, G + + +def get_gmm_output( + M: int, + N: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, +) -> Tensor: + assert M > 0, f"Number of out rows M must be positive (M = {M})." + assert N > 0, f"Number of out columns N must be positive (N = {N})." + + if existing_out is not None: + assert ( + existing_out.device == device + ), f"Existing output device and provided device don't match (existing = {existing_out.device}, provided = {device})." + assert ( + existing_out.dtype == preferred_element_type + ), f"Existing output type and preferred output type don't match (existing = {existing_out.dtype}, preferred = {preferred_element_type})." + assert existing_out.shape == ( + M, + N, + ), f"Existing output shape and GMM shape don't match (existing = {tuple(existing_out.shape)}, provided = {(M, N)})." + return existing_out + + return gen_gmm_output( + M, + N, + device=device, + preferred_element_type=preferred_element_type, + ) + + +def get_gmm_transposition(lhs: Tensor, rhs: Tensor, out: Tensor) -> tuple[bool, int]: + assert lhs.dim() == 2, f"lhs must have 2 dimensions (it's {lhs.dim()})." + assert rhs.dim() == 3, f"rhs must have 3 dimensions (it's {rhs.dim()})." + assert out.dim() == 2, f"out must have 2 dimensions (it's {out.dim()})." + + lhs_m, lhs_k = lhs.shape + G, rhs_k, rhs_n = rhs.shape + out_m, out_n = out.shape + + assert ( + lhs_m == out_m + ), f"M dimension of lhs and out don't match (lhs = {lhs_m}, rhs = {out_m})." + M = lhs_m + assert ( + lhs_k == rhs_k + ), f"K dimension of lhs and rhs don't match (lhs = {lhs_k}, rhs = {rhs_k})." + K = lhs_k + assert ( + rhs_n == out_n + ), f"N dimension of rhs and out don't match (lhs = {rhs_n}, rhs = {out_n})." + N = rhs_n + + assert M > 0, f"M must be positive, it's {M}." + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}" + assert G > 0, f"G must be positive, it's {G}" + + is_lhs_row_major = lhs.stride() == (K, 1) + assert is_lhs_row_major, "lhs must be row-major." + is_rhs_row_major = rhs.stride() == (K * N, N, 1) + is_rhs_col_major = rhs.stride() == (K * N, 1, K) + assert ( + is_rhs_row_major != is_rhs_col_major + ), "rhs must be row-major or column-major." + is_out_row_major = out.stride() == (N, 1) + assert is_out_row_major, "out must be row-major." + + # Get rhs leading dimension according to transposition configuration. + ld_rhs = N if is_rhs_row_major else K + + return is_rhs_col_major, ld_rhs + + +# TGMM helpers: tensor generation. +# ------------------------------------------------------------------------------ + + +def gen_tgmm_input( + M: int, + K: int, + N: int, + G: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, + trans_lhs: bool = TRANS_LHS, + rng_seed: int | None = RNG_SEED, + unif_group_sizes: bool = False, +) -> tuple[Tensor, Tensor, Tensor]: + assert K > 0, f"Number of lhs rows K must be positive (M = {K})." + assert M > 0, f"Number of lhs columns / rhs rows M must be positive (K = {M})." + assert N > 0, f"Number of rhs columns N must be positive (N = {N})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + if rng_seed is not None: + torch.manual_seed(rng_seed) + + if trans_lhs: + lhs = torch.randn((M, K), dtype=torch.float32, device=device).T + else: + lhs = torch.randn((K, M), dtype=torch.float32, device=device) + lhs = lhs.to(preferred_element_type) + + rhs = torch.randn((M, N), dtype=torch.float32, device=device) + rhs = rhs.to(preferred_element_type) + + group_sizes = ( + gen_uniform_group_sizes(M, G, device=device) + if unif_group_sizes + else gen_group_sizes(M, G, device=device, rng_seed=None) + ) + + return lhs, rhs, group_sizes + + +def gen_tgmm_output( + K: int, + N: int, + G: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, +) -> Tensor: + assert K > 0, f"Number of out rows K must be positive (K = {K})." + assert N > 0, f"Number of out columns N must be positive (N = {N})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + out = torch.empty((G, K, N), dtype=preferred_element_type, device=device) + + return out + + +def gen_tgmm_bias_grad( + K: int, + G: int, + device: torch.device | str = DEVICE, + with_bias_grad: bool = False, +) -> Tensor: + if with_bias_grad: + assert K > 0, f"Number of bias_grad rows K must be positive (K = {K})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + return torch.empty((G, K), device=device, dtype=torch.float32) + else: + # Return dummy pointer when bias_grad is not needed. + # Must be float32 because atomic_add does not support bf16/fp16, + # and Triton validates the pointer dtype even in dead branches. + return torch.tensor([], device=device, dtype=torch.float32) + + +def gen_tgmm_tensors( + M: int, + K: int, + N: int, + G: int, + num_group_sizes: int, + device: torch.device | str = DEVICE, + input_type: torch.dtype = DTYPE, + output_type: torch.dtype = DTYPE, + trans_lhs: bool = TRANS_LHS, + trans_rhs: bool = False, + rng_seed: int | None = RNG_SEED, + unif_group_sizes: bool = False, + use_bias: bool = False, +) -> tuple[Tensor, Tensor, list[Tensor], Tensor, Tensor | None]: + lhs, rhs, group_sizes_0 = gen_tgmm_input( + M, + K, + N, + G, + device=device, + preferred_element_type=input_type, + trans_lhs=trans_lhs, + rng_seed=rng_seed, + unif_group_sizes=unif_group_sizes, + ) + multiple_group_sizes = gen_multiple_group_sizes( + num_group_sizes, M, G, device=device, rng_seed=None, group_sizes_0=group_sizes_0 + ) + out = gen_tgmm_output(K, N, G, device=device, preferred_element_type=output_type) + if use_bias: + bias_grad = gen_tgmm_bias_grad(K, G, device=device, with_bias_grad=True) + else: + bias_grad = None + return lhs, rhs, multiple_group_sizes, out, bias_grad + + +# TGMM helpers: get information from tensors. +# ------------------------------------------------------------------------------ + + +def get_tgmm_shape( + lhs: Tensor, rhs: Tensor, group_sizes: Tensor +) -> tuple[int, int, int, int]: + assert lhs.dim() == 2, f"lhs must have 2 dimensions (it's {lhs.dim()})." + assert rhs.dim() == 2, f"rhs must have 2 dimensions (it's {rhs.dim()})." + assert ( + group_sizes.dim() == 1 + ), f"group_sizes must have 1 dimension (it's {group_sizes.dim()})." + + K, lhs_m = lhs.shape + rhs_m, N = rhs.shape + G = group_sizes.shape[0] + + assert ( + lhs_m == rhs_m + ), f"M dimension of lhs and rhs don't match (lhs = {lhs_m}, rhs = {rhs_m})." + M = lhs_m + + assert M > 0, f"M must be positive, it's {M}." + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}" + assert G > 0, f"G must be positive, it's {G}" + + return M, K, N, G + + +def get_tgmm_output( + K: int, + N: int, + G: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, +) -> Tensor: + assert K > 0, f"Number of out rows K must be positive (K = {K})." + assert N > 0, f"Number of out columns N must be positive (N = {N})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + if existing_out is not None: + assert ( + existing_out.device == device + ), f"Existing output device and provided device don't match (existing = {existing_out.device}, provided = {device})." + assert ( + existing_out.dtype == preferred_element_type + ), f"Existing output type and preferred output type don't match (existing = {existing_out.dtype}, preferred = {preferred_element_type})." + assert existing_out.shape == ( + G, + K, + N, + ), f"Existing output shape and GMM shape don't match (existing = {tuple(existing_out.shape)}, provided = {(G, K, N)})." + return existing_out + + return gen_tgmm_output( + K, + N, + G, + device=device, + preferred_element_type=preferred_element_type, + ) + + +def get_tgmm_bias_grad( + K: int, + G: int, + device: torch.device | str = DEVICE, + existing_bias_grad: Tensor | None = None, +) -> Tensor: + """ + Get or validate bias gradient tensor for TGMM. + + If existing_bias_grad is provided, validates its shape, device, dtype, and stride, + and always zeros it before returning (since the kernel uses atomic_add). + If existing_bias_grad is None, returns a dummy tensor (for use when COMPUTE_BIAS_GRAD=False). + Parameters + ---------- + K : int + Number of rows in the bias gradient tensor. + G : int + Number of groups. + device : torch.device or str + Device for the tensor. + existing_bias_grad : torch.Tensor or None + Existing bias gradient tensor to validate and use. + Returns + ------- + torch.Tensor + Valid bias gradient tensor or dummy tensor. + """ + assert K > 0, f"Number of bias_grad rows K must be positive (K = {K})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + if existing_bias_grad is not None: + # Validate existing bias_grad tensor. + expected_shape = (G, K) + assert ( + tuple(existing_bias_grad.shape) == expected_shape + ), f"bias_grad must have shape {expected_shape}, got {tuple(existing_bias_grad.shape)}." + assert ( + existing_bias_grad.device == device + ), f"bias_grad must be on the same device (bias_grad = {existing_bias_grad.device}, device = {device})." + assert ( + existing_bias_grad.dtype == torch.float32 + ), f"bias_grad must be torch.float32 (kernel uses atomic_add which requires float32), got {existing_bias_grad.dtype}." + assert existing_bias_grad.stride() == ( + K, + 1, + ), f"bias_grad must be row-major with stride (K, 1) = ({K}, 1), got {existing_bias_grad.stride()}." + + # Always zero the tensor since bias_grad represents gradients for the current + # computation and should start fresh. The kernel uses atomic_add which adds to + # existing values, so we must zero before the kernel runs. + existing_bias_grad.zero_() + + return existing_bias_grad + + else: + return gen_tgmm_bias_grad(K, G, device=device, with_bias_grad=False) + + +def get_tgmm_transposition(lhs: Tensor, rhs: Tensor, out: Tensor) -> tuple[bool, int]: + assert lhs.dim() == 2, f"lhs must have 2 dimensions (it's {lhs.dim()})." + assert rhs.dim() == 2, f"rhs must have 2 dimensions (it's {rhs.dim()})." + assert out.dim() == 3, f"out must have 3 dimensions (it's {out.dim()})." + + lhs_k, lhs_m = lhs.shape + rhs_m, rhs_n = rhs.shape + G, out_k, out_n = out.shape + + assert ( + lhs_m == rhs_m + ), f"M dimension of lhs and rhs don't match (lhs = {lhs_m}, rhs = {rhs_m})." + M = lhs_m + assert ( + lhs_k == out_k + ), f"K dimension of lhs and out don't match (lhs = {lhs_k}, rhs = {out_k})." + K = lhs_k + assert ( + rhs_n == out_n + ), f"N dimension of rhs and out don't match (lhs = {rhs_n}, rhs = {out_n})." + N = rhs_n + + assert M > 0, f"M must be positive, it's {M}." + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}" + assert G > 0, f"G must be positive, it's {G}" + + is_lhs_row_major = lhs.stride() == (M, 1) + is_lhs_col_major = lhs.stride() == (1, K) + assert ( + is_lhs_row_major != is_lhs_col_major + ), "lhs must be row-major or column-major." + is_rhs_row_major = rhs.stride() == (N, 1) + assert is_rhs_row_major, "rhs must be row-major." + is_out_row_major = out.stride() == (K * N, N, 1) + assert is_out_row_major, "out must be row-major." + + # Get lhs leading dimension according to transposition configuration. + ld_lhs = M if is_lhs_row_major else K + + return is_lhs_col_major, ld_lhs diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/logger.py b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..391ddf9b6543f5244e7f4932c8568d60748e15cd --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_grouped_gemm_triton/utils/logger.py @@ -0,0 +1,47 @@ +import os +import logging + + +# AITER Triton Logger which is singleton object around python logging. +# Note: Python logging is also a singleton object, but we want to read the +# env var AITER_LOG_LEVEL once at the beginning. Another alternative is to do +# this in __init__.py. In fact, that's how CK logger is setup. We can look at +# switching to that at some point +# +# AITER_LOG_LEVEL follows python logging levels +# DEBUG +# INFO +# WARNING +# ERROR +# CRITICAL +# +class AiterTritonLogger(object): + _instance = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super(AiterTritonLogger, cls).__new__(cls) + log_level_str = os.getenv("AITER_TRITON_LOG_LEVEL", "WARNING").upper() + numeric_level = getattr(logging, log_level_str, logging.WARNING) + cls._instance._logger = logging.getLogger("AITER_TRITON") + cls._instance._logger.setLevel(numeric_level) + + return cls._instance + + def get_logger(self): + return self._logger + + def debug(self, msg): + self._logger.debug(msg) + + def info(self, msg): + self._logger.info(msg) + + def warning(self, msg): + self._logger.warning(msg) + + def error(self, msg): + self._logger.error(msg) + + def critical(self, msg): + self._logger.critical(msg) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a720e7a2cc4e44636f6e433a2750e945dc38e8b2 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/__init__.py @@ -0,0 +1,10 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +# from megablocks.layers.dmoe import dMoE +from .moe import MoE + +__all__ = [ + 'MoE', + # 'dMoE', +] diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/activation_fn.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/activation_fn.py new file mode 100644 index 0000000000000000000000000000000000000000..0e1d956704840aa4daf7d1d71d24e051567feab9 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/activation_fn.py @@ -0,0 +1,33 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Callable, Union + +import torch +from ..stk import Matrix + + +def act_fn( + x: Matrix, + function: Callable, + return_grad_fn: bool = False, + **kwargs, +) -> Union[tuple[Matrix, Any] | Matrix]: + assert isinstance(x, Matrix) + with torch.set_grad_enabled(torch.is_grad_enabled() or return_grad_fn): + if return_grad_fn: + x.data.requires_grad = True + out = function(x.data, **kwargs) + y = Matrix( + x.size(), + out, + x.row_indices, + x.column_indices, + x.offsets, + x.column_indices_t, + x.offsets_t, + x.block_offsets_t, + ) + if return_grad_fn: + return y, out.backward + return y diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/all_to_all.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/all_to_all.py new file mode 100644 index 0000000000000000000000000000000000000000..5ac7067bcaa34db1d82b340c43550fe3577aa7a3 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/all_to_all.py @@ -0,0 +1,54 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.distributed as dist + + +class AllToAllOp(torch.autograd.Function): + + @staticmethod + def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op): + out = torch.empty((sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype) + + ctx.input_shape = x.shape + ctx.output_split_sizes = output_split_sizes + ctx.input_split_sizes = input_split_sizes + ctx.group = group + handle = dist.all_to_all_single( + out, + x, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group, + async_op=async_op, + ) + return out, handle + + @staticmethod + def backward(ctx, grad, _): + if ctx.needs_input_grad[0]: + out = torch.empty( + ctx.input_shape, + device=grad.device, + dtype=grad.dtype, + ) + dist.all_to_all_single( + out, + grad, + output_split_sizes=ctx.input_split_sizes, + input_split_sizes=ctx.output_split_sizes, + group=ctx.group, + ) + return out, None, None, None, None + return None, None, None, None, None + + +def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False): + return AllToAllOp.apply( + x, + output_split_sizes, + input_split_sizes, + group, + async_op, + ) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/arguments.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/arguments.py new file mode 100644 index 0000000000000000000000000000000000000000..4db9b1bd38bc2e2f421625c124f86b85f45c5ae0 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/arguments.py @@ -0,0 +1,101 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +from functools import partial +from typing import Any, Callable, Optional, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +# import megablocks.grouped_gemm_util as grouped_gemm +from .. import grouped_gemm_util as grouped_gemm + +# Type annotation for in-place Tensor initialization function. +InitFn = Union[Callable[[torch.Tensor], None], partial[torch.Tensor]] + +_ALLOWED_BITWIDTHS = (-1, 4, 8) + +DEFAULT_ACTIVATION_FN = partial(F.gelu, approximate='tanh') + + +@dataclasses.dataclass +class Arguments: + # Model arguments. + hidden_size: int = 1024 + ffn_hidden_size: int = 4096 + num_layers: int = 1 + bias: bool = True + return_bias: bool = True + activation_fn: Optional[Callable] = DEFAULT_ACTIVATION_FN + + # MoE arguments. + moe_num_experts: int = 1 + moe_top_k: int = 1 + moe_capacity_factor: int = 1 + moe_normalize_expert_weights: Optional[Union[int, float]] = None + moe_loss_weight: float = 0.1 + moe_jitter_eps: Optional[float] = None + moe_lbl_in_fp32: bool = False + + # Parallelism arguments. + moe_expert_model_parallelism: bool = False + expert_parallel_group: Optional[dist.ProcessGroup] = None + pipeline_model_parallel_size: int = 1 + num_layers_per_virtual_pipeline_stage: Optional[int] = None + + # Compute arguments. + memory_optimized_mlp: bool = False + mlp_type: str = 'mlp' + mlp_impl: str = 'sparse' + + # Initialization arguments. + fp16: bool = True + bf16: bool = False + device: Union[int, torch.device] = dataclasses.field(default_factory=torch.cuda.current_device) + init_method: InitFn = partial(torch.nn.init.normal_, mean=0.0, std=0.02) + output_layer_init_method: InitFn = init_method + + # Benchmarking arguments. + uniform_expert_assignment: bool = False + + # shared expert arguments + shared_expert: bool = False # enable using shared expert + fc_cls: Any = torch.nn.Linear # class of the fully connected layer in shared expert (purpose: to allow using custom FC layer eg te.Linear (for FP8)) + fc_kwargs: dict[str, Any] = dataclasses.field(default_factory=dict,) # kwargs for custom fc layers + remat_act_fn: bool = True # enable act fn to be rematerialized instead of stored + shared_expert_hidden_size: Optional[ + int] = None # hidden size of the shared expert IF we want to set it to something different from hidden_size + shared_expert_weighted_sum: bool = False # enable using weighted sum for shared expert output (wieghted by number of experts used) + + # Router Z-loss arguments + moe_zloss_weight: float = 0 # 1e-3 is a reasonable value + moe_zloss_in_fp32: bool = False + + def __post_init__(self): + # Sparse MLP is not supported with triton >=3.2.0 + # TODO: Remove this once sparse is supported with triton >=3.2.0 + if self.__getattribute__('mlp_impl') == 'sparse': + try: + import triton + if triton.__version__ >= '3.2.0': + raise ValueError( + 'Sparse MLP is not supported with triton >=3.2.0. Please use mlp_impl="grouped" instead.', + ) + except ImportError: + raise ImportError('Triton is required for sparse MLP implementation') + + if self.__getattribute__('mlp_impl') == 'grouped': + grouped_gemm.assert_grouped_gemm_is_available() + + if self.shared_expert_hidden_size is None: + self.shared_expert_hidden_size = self.ffn_hidden_size + + +def from_megatron(megatron_args: Any): + args = Arguments() + for field in dataclasses.fields(args): + if hasattr(megatron_args, field.name): + setattr(args, field.name, getattr(megatron_args, field.name)) + return args diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/common.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/common.py new file mode 100644 index 0000000000000000000000000000000000000000..2d07109702963ba48a3b94ab860807954dfd79c1 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/common.py @@ -0,0 +1,26 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from .arguments import Arguments + + +def dtype(args: Arguments): + if args.fp16: + return torch.float16 + elif args.bf16: + return torch.bfloat16 + return None + + +def cast_if_autocast_enabled(tensor): + if torch.is_autocast_enabled(): + if tensor.device.type == 'cuda': + dtype = torch.get_autocast_gpu_dtype() + elif tensor.device.type == 'cpu': + dtype = torch.get_autocast_cpu_dtype() + else: + raise NotImplementedError() + return tensor.to(dtype=dtype) + return tensor diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/dmlp_registry.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/dmlp_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..de2ed047042e438c7190ebb139b6f7f30009734c --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/dmlp_registry.py @@ -0,0 +1,42 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Union + +from . import glu, mlp +from .arguments import Arguments + +MlpType = Union[mlp.SparseMLP, glu.SparseGLU] + +_REGISTRY = { + 'mlp': { + 'grouped': mlp.GroupedMLP, + 'sparse': mlp.SparseMLP, + }, + 'glu': { + 'grouped': glu.GroupedGLU, + 'sparse': glu.SparseGLU, + }, +} + + +def get(args: Arguments) -> MlpType: + """Returns an MLP for use in a dMoE instance. + + Uses the provided arguments to instantiate the appropriate + MLP instance. This only contains MLPs for use in dMoEs + (ie. only for the dropless versions of MoEs). + + Args: + args: propagated Arguments dataclass. + + Returns: + An instantiated MLP constructed using the input args. + """ + if args.mlp_type not in _REGISTRY: + raise ValueError(f'Unsupported mlp type: {args.mlp_type}') + + if args.mlp_impl not in _REGISTRY[args.mlp_type]: + raise ValueError(f'{args.mlp_type} does not support {args.mlp_impl} backend.',) + + return _REGISTRY[args.mlp_type][args.mlp_impl](args) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/dmoe.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/dmoe.py new file mode 100644 index 0000000000000000000000000000000000000000..6d0375a4df2f27134c4127e60be04f3b45693050 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/dmoe.py @@ -0,0 +1,337 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import torch + +# try: +# import stk.ops +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/_layers/dmoe.py is needed.', +# ) + +# import megablocks.ops as ops +# # from megablocks.ops import ops +# from megablocks.layers import common, dmlp_registry, moe, mpu +# from megablocks.layers.arguments import Arguments + +from .. import stk +from .. import ops +from . import common, dmlp_registry, moe, mpu +from .arguments import Arguments + +def promote_scalar(x): + return x.view(1) if not len(x.size()) else x + + +class ParallelDroplessMLP(moe.ParallelMLP): + + def __init__(self, args: Arguments): + super(ParallelDroplessMLP, self).__init__(args) + self.hidden_size = args.hidden_size + self.ffn_hidden_size = mpu.features_per_rank(args) + self.blocking = 128 + self.mlp = dmlp_registry.get(args) + + # Calculate the number of bits needed to represent the column indices + # in the intermediate sparse matrix. + max_column_index = ((self.ffn_hidden_size * self.num_experts) // self.blocking) + self.transpose_sort_end_bit = max( + int(np.ceil(np.log2(max_column_index))), + 1, + ) + + def sparse_transpose(self, size, row_indices, column_indices, offsets): + block_columns = size[1] // self.blocking + + # Sort row indices by column indices to get the transposed matrix's + # column indices. + # + # NOTE: Our sort operation uses the same width indices as the input values. + # To avoid overflow when we have large activation matrices we cast to + # 32-bit before sorting. + _, gather_indices = ops.sort( + column_indices.int(), + self.transpose_sort_end_bit, + ) + + # There are a constant number of blocks in every row of the sparse matrix. + # A blocks offset is: + # + # row_index * blocks_per_row + column_index % blocks_per_row + # + # Once we have the block offsets ordered for transposition we can divide + # by blocks_per_row to get the transposed column indices. + column_indices_t = row_indices.gather(0, gather_indices.long()) + block_offsets_t = gather_indices.int() + + zero = torch.zeros((1,), dtype=torch.int32, device=row_indices.device) + nnz_per_column = ops.histogram(column_indices, block_columns) + nnz_per_column = ops.inclusive_cumsum(nnz_per_column, 0) + if nnz_per_column.dim() == 0: + # This addresses an edge case when ffn_hidden_size is equal to self.blocking. + nnz_per_column = nnz_per_column.unsqueeze(0) + offsets_t = torch.cat([zero, nnz_per_column]) + return column_indices_t, offsets_t, block_offsets_t + + def topology(self, x, padded_bins): + padded_tokens, _ = x.size() + assert padded_tokens % self.blocking == 0 + if self.ffn_hidden_size % self.blocking != 0: + raise ValueError( + f'The ffn_hidden_size {self.ffn_hidden_size} must be divisible by ' + + f'the block size {self.blocking}. Please update your configuration.', + ) + + # Offsets for the sparse matrix. All rows have the + # same number of nonzero blocks dictated by the + # dimensionality of a single expert. + block_rows = padded_tokens // self.blocking + blocks_per_row = self.ffn_hidden_size // self.blocking + offsets = torch.arange( + 0, + block_rows * blocks_per_row + 1, + blocks_per_row, + dtype=torch.int32, + device=x.device, + ) + + # Indices for the sparse matrix. The indices for + # the intermediate matrix are dynamic depending + # on the mapping of tokens to experts. + column_indices = ops.topology( + padded_bins, + self.blocking, + block_rows, + blocks_per_row, + ) + + # TODO(tgale): This is unused. Remove the need for this in stk. + # For now, use meta init to save the device memory. + data = torch.empty( + column_indices.numel(), + self.blocking, + self.blocking, + dtype=common.dtype(self.args), + device='meta', + ) + shape = ( + padded_tokens, + self.ffn_hidden_size * mpu.experts_per_rank(self.args), + ) + row_indices = stk.ops.row_indices(shape, data, offsets, column_indices) + column_indices_t, offsets_t, block_offsets_t = self.sparse_transpose( + shape, + row_indices, + column_indices, + offsets, + ) + return stk.Matrix( + shape, + data, + row_indices, + column_indices, + offsets, + column_indices_t, + offsets_t, + block_offsets_t, + ) + + def indices_and_padded_bins(self, top_experts): + # Sort the expert ids to produce the scatter/gather + # indices for the permutation. + top_experts = top_experts.int() + bin_ids, indices = ops.sort(top_experts, self.sort_end_bit) + + # Histogram the expert ids to identify the number of + # tokens routed to each expert. + tokens_per_expert = ops.histogram(top_experts, self.num_experts) + + # Round the token counts up to the block size used in + # the matrix muliplications. Caculate the starting + # position of each bin. + padded_tokens_per_expert = ops.round_up( + tokens_per_expert, + self.blocking, + ) + padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) + padded_bins = promote_scalar(padded_bins) + + # Calculate the bin bounds for the sorted tokens. + bins = ops.inclusive_cumsum(tokens_per_expert, 0) + bins = promote_scalar(bins) + return indices, bin_ids, bins, padded_bins, tokens_per_expert + + def sparse_forward_once(self, x, expert_weights, top_experts): + # x: [sl, bs, hs] + # expert_weights: [sl * bs, top-k] + # top_experts: [sl * bs, top-k] + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + with torch.no_grad(): + indices, bin_ids, bins, padded_bins, tokens_per_expert = (self.indices_and_padded_bins(top_experts)) + + # Route the tokens for MoE computation. + x = x.view(-1, x.shape[-1]) + x = ops.padded_gather( + x, + indices, + bin_ids, + bins, + padded_bins, + self.top_k, + ) + + # Create the sparse matrix topology. + with torch.no_grad(): + topo = self.topology(x, padded_bins) + + # Perform the expert computation. + x = self.mlp(x, topo) + + # Un-route the data for the MoE output. + x = ops.padded_scatter( + x, + indices, + bin_ids, + expert_weights, + bins, + padded_bins, + self.top_k, + ) + return x, tokens_per_expert + + # For use in the base-class parallel_forward_once. + def sparse_permute_and_compute( + self, + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, # unused + top_k, + ): + + # Round the token counts up to the block size used in the matrix + # multiplication. Calculate the starting position of each bin. + padded_tokens_per_expert = ops.round_up( + tokens_per_expert, + self.blocking, + ) + padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) + padded_bins = promote_scalar(padded_bins) + + # Route the tokens for MoE computation. + x = x.view(-1, x.shape[-1]) + x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k) + + # Create the sparse matrix topology. + with torch.no_grad(): + topo = self.topology(x, padded_bins) + + # Perform the expert computation. + x = self.mlp(x, topo) + + # Un-route the data for the MoE output. + return ops.padded_scatter( + x, + indices, + bin_ids, + expert_weights, + bins, + padded_bins, + top_k, + ) + + def grouped_forward_once(self, x, expert_weights, top_experts): + # x: [sl, bs, hs] + # expert_weights: [sl * bs, top-k] + # top_experts: [sl * bs, top-k] + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + with torch.no_grad(): + indices, bin_ids, bins, tokens_per_expert = (self.indices_and_bins(top_experts)) + + out = self.grouped_permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + -1, # unused + self.args.moe_top_k, + ) + return out, tokens_per_expert + + def grouped_permute_and_compute( + self, + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, # unused + top_k, + ): + + # Route the tokens for MoE computation. + x = x.view(-1, x.shape[-1]) + x = ops.gather(x, indices, bin_ids, bins, top_k) + + # Perform the expert computation. + x = self.mlp(x, tokens_per_expert) + + # Un-route the data for the MoE output. + return ops.scatter(x, indices, bin_ids, expert_weights, bins, top_k) + + def forward_once(self, x, expert_weights, top_experts): + if self.args.mlp_impl == 'sparse': + return self.sparse_forward_once(x, expert_weights, top_experts) + else: + return self.grouped_forward_once(x, expert_weights, top_experts) + + def permute_and_compute( + self, + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, + top_k, + ): + if self.args.mlp_impl == 'sparse': + return self.sparse_permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, + top_k, + ) + else: + return self.grouped_permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, + top_k, + ) + + +class dMoE(moe.MoE): + + def _init_experts_mlp(self, args: Arguments): + return ParallelDroplessMLP(args) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/gelu.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/gelu.py new file mode 100644 index 0000000000000000000000000000000000000000..c4c9e6532798615b5c12c96694241a4c18ee8f7b --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/gelu.py @@ -0,0 +1,52 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +# try: +# import stk +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/_layers/gelu.py is needed.', +# ) + +from .. import stk + +import torch +import torch.nn.functional as F + + +@torch.jit.script +def _gelu_backward_inplace(g, x): + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + ff = (0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)) + return g.mul_(ff) + + +def gelu_backward_(grad: stk.Matrix, x: stk.Matrix): + # NOTE: The two sparse matrices must have the same topology. + if isinstance(grad, stk.Matrix) and isinstance(x, stk.Matrix): + return stk.Matrix( + x.size(), + _gelu_backward_inplace(grad.data, x.data), + x.row_indices, + x.column_indices, + x.offsets, + x.column_indices_t, + x.offsets_t, + x.block_offsets_t, + ) + return _gelu_backward_inplace(grad, x) + + +def gelu(x: stk.Matrix): + assert isinstance(x, stk.Matrix) + return stk.Matrix( + x.size(), + F.gelu(x.data, approximate='tanh'), + x.row_indices, + x.column_indices, + x.offsets, + x.column_indices_t, + x.offsets_t, + x.block_offsets_t, + ) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/glu.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/glu.py new file mode 100644 index 0000000000000000000000000000000000000000..5f297a41ff6a1a2a285f5b461951672364b898da --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/glu.py @@ -0,0 +1,244 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +# import stk.ops +# try: +# import stk.ops +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/_layers/glu.py is needed.', +# ) + +from .. import stk + +import torch + +# from megablocks import grouped_gemm_util as gg +# from megablocks.layers import common, mpu +# from megablocks.layers.activation_fn import act_fn +# from megablocks.layers.arguments import Arguments +# from megablocks.layers.mlp import ( +# SharedMLP, +# SparseMLP, +# create_dmoe_expert_weights, +# resolve_dtensor, +# ) + +from .. import grouped_gemm_util as gg +from . import common, mpu +from .activation_fn import act_fn +from .arguments import Arguments +from .mlp import ( + SharedMLP, + SparseMLP, + create_dmoe_expert_weights, + resolve_dtensor, +) + + +class SparseGLU(SparseMLP): + + def __init__(self, args: Arguments): + super().__init__(args) + self.v1 = torch.nn.Parameter( + torch.empty( + self._num_rows_per_rank, + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + with torch.no_grad(): + self.v1.copy_( + create_dmoe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.init_method, + ), + ) + + mpu.set_expert_model_parallel_attributes( + self.v1, + self._should_set_parallelism_attribute, + ) + + def forward(self, x, topo): + if self.args.memory_optimized_mlp: + raise NotImplementedError( + 'Memory optimized implementation not yet supported with GLU with sparse kernels.', + ) + + w1, v1, w2 = self.scale_grad(self.w1), self.scale_grad(self.v1,), self.scale_grad(self.w2) + w1, v1, w2 = resolve_dtensor(w1), resolve_dtensor(v1,), resolve_dtensor(w2) + + # Compute the GLU. + x1 = stk.ops.sdd(x, w1.t(), topo) + x2 = stk.ops.sdd(x, v1.t(), topo) + + activation_fn_out = act_fn(x1, self.args.activation_fn) + x1 = stk.ops.mul(activation_fn_out, x2) + + return stk.ops.dsd(x1, w2) + + +class MemoryOptimizedGroupedGLU(torch.autograd.Function): + """GroupedMLP with manually scheduled memory reuse.""" + + @staticmethod + @torch.amp.autocast_mode.custom_fwd(device_type='cuda') + def forward(ctx, x, w1, v1, w2, batch_sizes, activation_fn): + # Cast inputs using ctx dtype from AMP + if ctx._fwd_used_autocast: + x = x.to(ctx._dtype) + w1 = w1.to(ctx._dtype) + v1 = v1.to(ctx._dtype) + w2 = w2.to(ctx._dtype) + # x: [m, k], w1: [n, k], v1: [n, k], w2: [n, k] + if (not x.is_contiguous() or not w1.is_contiguous() or not v1.is_contiguous() or not w2.is_contiguous()): + raise ValueError("Expected contiguous 'x', 'w1', 'v1' and 'w2'.") + + # Layer 0: x @ w1.t(). + assert gg.backend is not None + sdd_out = gg.backend.gmm(x, w1, batch_sizes, trans_b=True) + v1_out = gg.backend.gmm(x, v1, batch_sizes, trans_b=True) + + # GeLU. + activation_fn_out = activation_fn(sdd_out) * v1_out + + # Layer 1: x @ w2. + dsd_out = gg.backend.gmm(activation_fn_out, w2, batch_sizes) + + # NOTE: Save the input to the layer and the activation_fn input for + # gradient computation. We'll re-compute the activation_fn forward + # pass in the backward pass to avoid materializing another + # intermediate. + ctx.x_shape = x.shape + ctx.sdd_out_shape = sdd_out.shape + ctx.dtype = x.dtype + ctx.activation_fn = activation_fn + ctx.save_for_backward(w1, v1, w2, batch_sizes, x, sdd_out, v1_out) + return dsd_out + + @staticmethod + @torch.amp.autocast_mode.custom_bwd(device_type='cuda') + def backward(ctx, ddsd_out): + if (not ctx.needs_input_grad[0] or not ctx.needs_input_grad[1] or not ctx.needs_input_grad[2]): + raise ValueError('Expected all MLP inputs to need grad.') + + # Unpack saved tensors + # dtype = ctx.dtype + saved_tensors = ctx.saved_tensors + w1, v1, w2 = saved_tensors[:3] + batch_sizes = saved_tensors[3] + x = saved_tensors[4] + sdd_out, v1_out = saved_tensors[5:7] + + # Rematerialize activation_fn output. + activation_fn = ctx.activation_fn + with torch.set_grad_enabled(True): + sdd_out.requires_grad = True + v1_out.requires_grad = True + activation_fn_out = activation_fn(sdd_out) * v1_out + activation_grad_fn = activation_fn_out.backward + + # Compute dw2 with recomputed activation_fn output. + assert gg.backend is not None + dw2 = gg.backend.gmm( + activation_fn_out, + ddsd_out, + batch_sizes, + trans_a=True, + ) + + # Compute dactivation_fn_out. + # + # NOTE: We reuse the activation_fn_out allocation. + dactivation_fn_out = activation_fn_out + gg.backend.gmm( + ddsd_out, + w2, + batch_sizes, + trans_b=True, + c=dactivation_fn_out, + ) + + # Compute dsdd_out. + # + # NOTE: This reuses the dactivation_fn_out allocation. + assert activation_grad_fn is not None + activation_grad_fn(dactivation_fn_out) + dsdd_out = sdd_out.grad + dv1_out = v1_out.grad + + # Compute dw1. + dw1 = gg.backend.gmm(dsdd_out, x, batch_sizes, trans_a=True) + + # Compute dv1. + dv1 = gg.backend.gmm(dv1_out, x, batch_sizes, trans_a=True) + + # Compute dx. + # + # NOTE: This reuses the ddsd_out allocation. + dx = ddsd_out + gg.backend.gmm(dsdd_out, w1, batch_sizes, c=dx) + dx += gg.backend.gmm(dv1_out, v1, batch_sizes) + return dx, dw1, dv1, dw2, None, None + + +memory_optimized_grouped_glu = MemoryOptimizedGroupedGLU.apply + + +class GroupedGLU(SparseGLU): + + def forward(self, x, tokens_per_expert): + batch_sizes = tokens_per_expert.cpu().to(torch.long) + w1, v1, w2 = ( + self.scale_grad(self.w1), + self.scale_grad(self.v1), + self.scale_grad(self.w2), + ) + w1, v1, w2 = resolve_dtensor(w1), resolve_dtensor(v1,), resolve_dtensor(w2) + + # Re-shape the weights for the grouped GEMMs. + ne = mpu.experts_per_rank(self.args) + w1 = w1.view(ne, -1, self.args.hidden_size) + v1 = v1.view(ne, -1, self.args.hidden_size) + w2 = w2.view(ne, -1, self.args.hidden_size) + + if self.args.memory_optimized_mlp: + return memory_optimized_grouped_glu( + x, + w1, + v1, + w2, + batch_sizes, + self.args.activation_fn, + ) + + # Compute the MLP. + assert gg.ops is not None + x1 = gg.ops.gmm(x, w1, batch_sizes, trans_b=True) + x2 = gg.ops.gmm(x, v1, batch_sizes, trans_b=True) + x1 = self.args.activation_fn(x1) * x2 + return gg.ops.gmm(x1, w2, batch_sizes) + + +class SharedGLU(SharedMLP): + """GPU for shared expert. + + Note: this is a copy -> pasta -> modify of the LLM-Foundry MPTGLU class + """ + + def __init__(self, args: Arguments): + super().__init__(args) + self.gate_proj = args.fc_cls( + args.hidden_size, + self.args.shared_expert_hidden_size, + **self.fc_kwargs, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x)) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/memory_test.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/memory_test.py new file mode 100644 index 0000000000000000000000000000000000000000..74d1166931b712635131985b25a89f4ca23e576d --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/memory_test.py @@ -0,0 +1,103 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import gc + +import torch +import torch.distributed as dist + +# from megablocks.layers import arguments, dmoe +from . import arguments, dmoe + +_TESTS = ((8, 2048, 4096, 4096, 32, 4),) + + +def get_tensors(): + ptrs = set() + out = [] + for obj in gc.get_objects(): + if torch.is_tensor(obj): + if not obj.is_contiguous() or obj.data_ptr() in ptrs: + continue + out.append(obj) + ptrs.add(obj.data_ptr()) + return out + + +def test_memory( + group, + batch_size, + sequence_length, + hidden_size, + ffn_hidden_size, + num_experts, + top_k, +): + args = arguments.Arguments( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + moe_num_experts=num_experts, + moe_top_k=top_k, + moe_expert_model_parallelism=True, + expert_parallel_group=group, + fp16=False, + bf16=True, + device=torch.cuda.current_device(), + ) + layer = dmoe.dMoE(args).cuda() + + x = torch.randn((batch_size, sequence_length, hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16).requires_grad_(True) + torch.cuda.empty_cache() + + # Run forward + backward. + # with torch.autograd.detect_anomaly(): + out, _ = layer(x) + out.mean().backward() + + # Report peak memory. + mem = torch.cuda.max_memory_allocated() + print('Max Memory Allocated = {:0.0f}MiB'.format(mem / 1e6)) + print('Max Memory Reserved = {:0.0f}MiB'.format(torch.cuda.max_memory_reserved() / 1e6,),) + + # Calculate weight and gradient memory usage. + weight_memory = 2 * ( + layer.router.layer.weight.numel() + layer.experts.mlp.w1.numel() + layer.experts.mlp.w2.numel() + ) + + def grad_numel(x): + if x.grad is not None: + return x.grad.numel() + return 0 + + grad_memory = 2 * ( + grad_numel(layer.router.layer.weight) + grad_numel(layer.experts.mlp.w1) + grad_numel(layer.experts.mlp.w2) + ) + weight_memory += grad_memory + + print('Weight Memory Allocated = {:0.0f}MiB'.format(weight_memory / 1e6)) + print('Activation Memory Allocated = {:0.0f}MiB'.format((mem - weight_memory) / 1e6,),) + + # Manually calculate GPU memory usage from the garbage + # collector. + gc.collect() + total = 0 + tensors = get_tensors() + tensors = sorted(tensors, key=lambda x: -x.numel()) + for i, t in enumerate(tensors): + total += t.numel() + print(f'{i}: {t.shape}, {t.numel() * 2}') + del tensors + + print('Total Bytes Found = {:0.0f}MiB'.format(total * 2 / 1e6)) + + +if __name__ == '__main__': + assert dist.is_available() + group = dist.init_process_group(backend='nccl') + local_rank = dist.get_rank(group) + torch.cuda.set_device(local_rank) + + for args in _TESTS: + test_memory(group, *args) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/mlp.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..c99afb9904c24a8b6a83e79059cd1251dbbfd99e --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/mlp.py @@ -0,0 +1,587 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# try: +# import stk +# import stk.backend.triton_kernels +# import stk.ops +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/_layers/mlp.py is needed.', +# ) + +from .. import stk + +import torch +from packaging import version + +# from megablocks import grouped_gemm_util as gg +# from megablocks.layers import common, gelu, mpu +# from megablocks.layers.activation_fn import act_fn +# from megablocks.layers.arguments import DEFAULT_ACTIVATION_FN, Arguments, InitFn + +from .. import grouped_gemm_util as gg +from . import common, gelu, mpu +from .activation_fn import act_fn +from .arguments import DEFAULT_ACTIVATION_FN, Arguments, InitFn + +class ScaleGradient(torch.autograd.Function): + + @staticmethod + @torch.amp.autocast_mode.custom_fwd(device_type='cuda') + def forward(ctx: Any, x: torch.Tensor, scale: float): + ctx.scale = scale + return x + + @staticmethod + @torch.amp.autocast_mode.custom_bwd(device_type='cuda') + def backward(ctx: torch.Tensor, grad: torch.Tensor): + return grad * ctx.scale, None + + +scale_gradient = ScaleGradient.apply + + +def resolve_dtensor(weight: torch.Tensor): + if version.parse(torch.__version__) >= version.parse('2.0.0'): + from torch.distributed._tensor import DTensor + if isinstance(weight, DTensor): + return weight.to_local() + return weight + + +def create_moe_expert_weights( + args: Arguments, + num_experts: int, + ffn_hidden_size: int, + hidden_size: int, + init_method: InitFn, +): + # Create the entire weight matrix such that the sampled weights will + # not vary between data parallelism and expert model parallelism for + # the same random seed. + master_weights = torch.empty( + num_experts, + ffn_hidden_size, + hidden_size, + device=args.device, + dtype=common.dtype(args), + ) + init_method(master_weights) + + if not args.moe_expert_model_parallelism: + return master_weights + + # Calculate the amount of sharding in each dimension. + expert_sharding_degree = mpu.expert_sharding_degree(args) + hidden_sharding_degree = mpu.hidden_sharding_degree(args) + + # Calculate the experts per rank. + # + # NOTE: We assign ranks to be expert parallel before going + # tensor parallel. + rank = mpu.get_expert_parallel_rank(args) + expert_rank = rank % expert_sharding_degree + num_experts_per_rank = num_experts // expert_sharding_degree + start_expert = expert_rank * num_experts_per_rank + end_expert = (expert_rank + 1) * num_experts_per_rank + + # Calculate the rows per rank. + row_rank = rank // expert_sharding_degree + num_rows_per_rank = ffn_hidden_size // hidden_sharding_degree + start_row = row_rank * num_rows_per_rank + end_row = (row_rank + 1) * num_rows_per_rank + + # Slice the weight matrix to get the chunk for this rank. + with torch.no_grad(): + weights = master_weights[start_expert:end_expert, start_row:end_row] + return weights + + +class MLP(torch.nn.Module): + + def __init__(self, args: Arguments): + super().__init__() + self.args = args + # expert_parallel_world_size = mpu.get_expert_parallel_world_size(args) + experts_per_rank = mpu.experts_per_rank(args) + + self.w1 = torch.nn.Parameter( + torch.empty( + experts_per_rank, + args.hidden_size, + mpu.features_per_rank(args), + device=args.device, + dtype=common.dtype(args), + ), + ) + self.w2 = torch.nn.Parameter( + torch.empty( + experts_per_rank, + mpu.features_per_rank(args), + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + mpu.set_expert_model_parallel_attributes( + self.w1, + args.moe_expert_model_parallelism, + ) + mpu.set_expert_model_parallel_attributes( + self.w2, + args.moe_expert_model_parallelism, + ) + + # Initialize the parameters for the MLP. + # + # NOTE: It is important that we create the weight tensors prior + # to creating the master weights and slicing our the piece for + # this rank. If the master weights are created first the PyTorch + # caching allocator appears to use the same memory block for these + # and the slice which causes large increases in our peak memory + # usage. + with torch.no_grad(): + w1 = create_moe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.init_method, + ) + self.w1.copy_(w1.transpose(1, 2).contiguous()) + self.w2.copy_( + create_moe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.output_layer_init_method, + ), + ) + + self.gradient_scale = None + if self.args.moe_expert_model_parallelism: + self.gradient_scale = 1 / mpu.get_expert_parallel_world_size(self.args,) + + def scale_grad(self, w): + if self.gradient_scale is None: + return w + return scale_gradient(w, self.gradient_scale) + + def forward(self, x): + w1, w2 = self.scale_grad(self.w1), self.scale_grad(self.w2) + w1, w2 = resolve_dtensor(w1), resolve_dtensor(w2) + x = torch.bmm(x, w1) + x = self.args.activation_fn(x) + return torch.bmm(x, w2) + + +def create_dmoe_expert_weights( + args: Arguments, + num_experts: int, + rows: int, + columns: int, + init_method: InitFn, +): + weights = create_moe_expert_weights( + args, + num_experts, + rows, + columns, + init_method, + ) + return weights.view([-1, columns]) + + +class MemoryOptimizedMLP(torch.autograd.Function): + """Sparse MLP with manually scheduled memory reuse.""" + + @staticmethod + @torch.amp.autocast_mode.custom_fwd(device_type='cuda') + def forward(ctx, x, w1, w2, topo, activation_fn): + # Cast inputs using ctx dtype from AMP + if ctx._fwd_used_autocast: + x = x.to(ctx._dtype) + w1 = w1.to(ctx._dtype) + w2 = w2.to(ctx._dtype) + # x: [m, k], w1: [n, k], w2: [n, k] + if (not x.is_contiguous() or not w1.is_contiguous() or not w2.is_contiguous()): + raise ValueError("Expected contiguous 'x', 'w1' and 'w2'.") + + topo_tensors = ( + topo.row_indices, + topo.column_indices, + topo.offsets, + topo.column_indices_t, + topo.offsets_t, + topo.block_offsets_t, + ) + + # Layer 0: x @ w1.t(). + sdd_out = stk.ops.sdd(x, w1.t(), topo) + + # GeLU. + activation_fn_out = act_fn(sdd_out, activation_fn) + + # Layer 1: x @ w2. + dsd_out = stk.ops.dsd(activation_fn_out, w2) + + # NOTE: Save the input to the layer and the activation_fn input for + # gradient computation. We'll re-compute the activation_fn forward + # pass in the backward pass to avoid materializing another + # intermediate. + ctx.shape = topo.shape + ctx.x_shape = x.shape + ctx.sdd_out_shape = sdd_out.data.shape + ctx.dtype = x.dtype + ctx.activation_fn = activation_fn + ctx.save_for_backward(w1, w2, *topo_tensors, x, sdd_out.data) + return dsd_out + + @staticmethod + @torch.amp.autocast_mode.custom_bwd(device_type='cuda') + def backward(ctx, ddsd_out): + if (not ctx.needs_input_grad[0] or not ctx.needs_input_grad[1] or not ctx.needs_input_grad[2]): + raise ValueError('Expected all MLP inputs to need grad.') + + # unpack saved tensors + # dtype = ctx.dtype + saved_tensors = ctx.saved_tensors + w1, w2 = saved_tensors[:2] + topo_tensors = saved_tensors[2:8] + x = saved_tensors[8] + sdd_out_data = saved_tensors[9] + + # rematerialize activation function output + activation_fn = ctx.activation_fn + sdd_out = stk.Matrix(ctx.shape, sdd_out_data, *topo_tensors) + activation_fn_out, activation_grad_fn = act_fn( + sdd_out, + activation_fn, + return_grad_fn=True, + ) + + # Compute dw2 with recomputed activation_fn output. + dw2 = stk.ops.dsd(activation_fn_out.t(), ddsd_out) + + # Compute dactivation_fn_out. + # + # NOTE: We reuse the activation_fn_out allocation. + dactivation_fn_out = activation_fn_out + stk.backend.triton_kernels.sdd( + ddsd_out, + w2.t(), + dactivation_fn_out.shape, + dactivation_fn_out.data, + dactivation_fn_out.offsets, + dactivation_fn_out.row_indices, + dactivation_fn_out.column_indices, + ) + + # Compute dsdd_out. + # + # NOTE: This reuses the dactivation_fn_out allocation. + if activation_fn is DEFAULT_ACTIVATION_FN: + dsdd_out = gelu.gelu_backward_(dactivation_fn_out, sdd_out) + else: + assert activation_grad_fn is not None + activation_grad_fn(dactivation_fn_out.data) + dsdd_out = stk.Matrix(ctx.shape, sdd_out.data.grad, *topo_tensors) + + # Compute dw1. + dw1 = stk.ops.dsd(dsdd_out.t(), x) + + # Compute dx. + # + # NOTE: This reuses the ddsd_out allocation. + stk.backend.triton_kernels.dsd( + dsdd_out.shape, + dsdd_out.data, + dsdd_out.offsets, + dsdd_out.row_indices, + dsdd_out.column_indices, + dsdd_out.offsets_t, + dsdd_out.column_indices_t, + dsdd_out.block_offsets_t, + False, + w1, + ddsd_out, + ) + dx = ddsd_out + return dx, dw1, dw2, None, None + + +memory_optimized_mlp = MemoryOptimizedMLP.apply + + +class SparseMLP(torch.nn.Module): + + def __init__(self, args: Arguments): + super().__init__() + self.args = args + self._num_rows_per_rank = mpu.experts_per_rank(args) * mpu.features_per_rank(args) + + self.w1 = torch.nn.Parameter( + torch.empty( + self._num_rows_per_rank, + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + self.w2 = torch.nn.Parameter( + torch.empty( + self._num_rows_per_rank, + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + + # Initialize the parameters for the MLP. + # + # NOTE: It is important that we create the weight tensors prior + # to creating the master weights and slicing our the piece for + # this rank. If the master weights are created first the PyTorch + # caching allocator appears to use the same memory block for these + # and the slice which causes large increases in our peak memory + # usage. + with torch.no_grad(): + self.w1.copy_( + create_dmoe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.init_method, + ), + ) + self.w2.copy_( + create_dmoe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.output_layer_init_method, + ), + ) + + self._should_set_parallelism_attribute = args.moe_expert_model_parallelism + mpu.set_expert_model_parallel_attributes( + self.w1, + self._should_set_parallelism_attribute, + ) + mpu.set_expert_model_parallel_attributes( + self.w2, + self._should_set_parallelism_attribute, + ) + + self.gradient_scale = None + if self.args.moe_expert_model_parallelism: + self.gradient_scale = 1 / mpu.get_expert_parallel_world_size(self.args,) + + def scale_grad(self, w): + if self.gradient_scale is None: + return w + return scale_gradient(w, self.gradient_scale) + + def forward(self, x, topo): + w1, w2 = self.scale_grad(self.w1), self.scale_grad(self.w2) + w1, w2 = resolve_dtensor(w1), resolve_dtensor(w2) + if self.args.memory_optimized_mlp: + return memory_optimized_mlp( + x, + w1, + w2, + topo, + self.args.activation_fn, + ) + + # Compute the MLP. + x = stk.ops.sdd(x, w1.t(), topo) + activation_fn_out = act_fn(x, self.args.activation_fn) + return stk.ops.dsd(activation_fn_out, w2) + + +class MemoryOptimizedGroupedMLP(torch.autograd.Function): + """GroupedMLP with manually scheduled memory reuse.""" + + @staticmethod + @torch.amp.autocast_mode.custom_fwd(device_type='cuda') + def forward(ctx, x, w1, w2, batch_sizes, activation_fn): + # Cast inputs using ctx dtype from AMP + if ctx._fwd_used_autocast: + x = x.to(ctx._dtype) + w1 = w1.to(ctx._dtype) + w2 = w2.to(ctx._dtype) + # x: [m, k], w1: [n, k], w2: [n, k] + if (not x.is_contiguous() or not w1.is_contiguous() or not w2.is_contiguous()): + raise ValueError("Expected contiguous 'x', 'w1' and 'w2'.") + + # Layer 0: x @ w1.t(). + assert gg.backend is not None + sdd_out = gg.backend.gmm(x, w1, batch_sizes, trans_b=True) + + # activation_fn + activation_fn_out = activation_fn(sdd_out) + + # Layer 1: x @ w2. + dsd_out = gg.backend.gmm(activation_fn_out, w2, batch_sizes) + + # NOTE: Save the input to the layer and the activation_fn input for + # gradient computation. We'll re-compute the activation_fn forward + # pass in the backward pass to avoid materializing another + # intermediate. + ctx.x_shape = x.shape + ctx.sdd_out_shape = sdd_out.shape + ctx.dtype = x.dtype + ctx.activation_fn = activation_fn + ctx.save_for_backward(w1, w2, batch_sizes, x, sdd_out) + return dsd_out + + @staticmethod + @torch.amp.autocast_mode.custom_bwd(device_type='cuda') + def backward(ctx: Any, ddsd_out: torch.Tensor): + if (not ctx.needs_input_grad[0] or not ctx.needs_input_grad[1] or not ctx.needs_input_grad[2]): + raise ValueError('Expected all MLP inputs to need grad.') + + # Unpack saved tensors + # dtype = ctx.dtype + saved_tensors = ctx.saved_tensors + w1, w2 = saved_tensors[:2] + batch_sizes = saved_tensors[2] + x = saved_tensors[3] + sdd_out = saved_tensors[4] + + # Rematerialize activation_fn output. + activation_fn = ctx.activation_fn + with torch.set_grad_enabled(True): + sdd_out.requires_grad = True + activation_fn_out = activation_fn(sdd_out) + activation_grad_fn = activation_fn_out.backward + + # Compute dw2 with recomputed activation_fn output. + assert gg.backend is not None + dw2 = gg.backend.gmm( + activation_fn_out, + ddsd_out, + batch_sizes, + trans_a=True, + ) + + # Compute dactivation_fn_out. + # + # NOTE: We reuse the activation_fn_out allocation. + dactivation_fn_out = activation_fn_out + gg.backend.gmm( + ddsd_out, + w2, + batch_sizes, + trans_b=True, + c=dactivation_fn_out, + ) + + # Compute dsdd_out. + # + # NOTE: This reuses the dactivation_fn_out allocation. + if activation_fn is DEFAULT_ACTIVATION_FN: + dsdd_out = gelu.gelu_backward_(dactivation_fn_out, sdd_out) + else: + assert activation_grad_fn is not None + activation_grad_fn(dactivation_fn_out) + dsdd_out = sdd_out.grad + + # Compute dw1. + dw1 = gg.backend.gmm(dsdd_out, x, batch_sizes, trans_a=True) + + # Compute dx. + # + # NOTE: This reuses the ddsd_out allocation. + gg.backend.gmm(dsdd_out, w1, batch_sizes, c=ddsd_out) + dx = ddsd_out + return dx, dw1, dw2, None, None + + +memory_optimized_grouped_mlp = MemoryOptimizedGroupedMLP.apply + + +class GroupedMLP(SparseMLP): + + def forward(self, x, tokens_per_expert): + batch_sizes = tokens_per_expert.cpu().to(torch.long) + w1, w2 = (self.scale_grad(self.w1), self.scale_grad(self.w2)) + + # Re-shape the weights for the grouped GEMMs. + ne = mpu.experts_per_rank(self.args) + w1 = resolve_dtensor(w1).view(ne, -1, self.args.hidden_size) + w2 = resolve_dtensor(w2).view(ne, -1, self.args.hidden_size) + + if self.args.memory_optimized_mlp: + return memory_optimized_grouped_mlp( + x, + w1, + w2, + batch_sizes, + self.args.activation_fn, + ) + + # Compute the MLP. + assert gg.ops is not None + x = gg.ops.gmm(x, w1, batch_sizes, trans_b=True) + x = self.args.activation_fn(x) + return gg.ops.gmm(x, w2, batch_sizes) + + +class SharedMLP(torch.nn.Module): + """MLP for shared expert. + + Note: this is a copy -> pasta -> modify of the LLM-Foundry MPTMLP class + """ + + def __init__(self, args: Arguments): + super().__init__() + self.args = args + self.fc_kwargs: dict[str, Any] = { + 'bias': args.bias, + 'device': args.device, + } + self.fc_kwargs.update(args.fc_kwargs) + + self.up_proj = args.fc_cls( + args.hidden_size, + args.shared_expert_hidden_size, + **self.fc_kwargs, + ) + self.act = args.activation_fn + self.down_proj = args.fc_cls( + args.shared_expert_hidden_size, + args.hidden_size, + **self.fc_kwargs, + ) + self.down_proj._is_residual = True # a flag for llm-foundry init + + def add_experts_sharedexpert( + self, + shared_expert_out: torch.Tensor, + expert_out: torch.Tensor, + ) -> torch.Tensor: + # Helper function to add expert output to shared expert output + # with optional weighted sum. + if self.args.shared_expert_weighted_sum: + # enable using weighted sum for shared expert output + # wieghted by number of experts used + t_experts = self.args.moe_top_k + 1 + sh_mlp_out = shared_expert_out / t_experts + return sh_mlp_out.add( + expert_out, + alpha=(self.args.moe_top_k / t_experts), + ) + + return shared_expert_out + expert_out + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act(self.up_proj(x))) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/moe.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/moe.py new file mode 100644 index 0000000000000000000000000000000000000000..d0a4aeaacc9c86fc70944e730c53f7a55644e05e --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/moe.py @@ -0,0 +1,507 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Optional, Tuple + +import numpy as np +import torch +import torch.distributed as dist + +# import megablocks.ops as ops +# from megablocks.layers import common, mlp, mpu, router, sharedexpert_registry +# from megablocks.layers.all_to_all import all_to_all +# from megablocks.layers.arguments import Arguments + +from ..ops import ( + sort, + histogram, + inclusive_cumsum, + exclusive_cumsum, + binned_gather, + binned_scatter, + gather, + scatter, + repeat, + replicate, +) + +from . import common, mlp, mpu, router, sharedexpert_registry +from .arguments import Arguments +from .all_to_all import all_to_all + +_LOAD_BALANCING_LOSS = [] + + +def save_load_balancing_loss(loss): + global _LOAD_BALANCING_LOSS + _LOAD_BALANCING_LOSS.append(loss) + + +def get_load_balancing_loss(): + global _LOAD_BALANCING_LOSS + return _LOAD_BALANCING_LOSS + + +def clear_load_balancing_loss(): + global _LOAD_BALANCING_LOSS + _LOAD_BALANCING_LOSS.clear() + + +def batched_load_balancing_loss(args: Arguments): + if args.moe_loss_weight == 0: + return 0.0 + + # tokens_per_expert[i].shape = (num_experts) + # expert_scores[i].shape = (tokens, num_experts) + tokens_per_expert, expert_scores = zip(*get_load_balancing_loss()) + num_layers_per_pipeline_stage = (args.num_layers // args.pipeline_model_parallel_size) + if args.num_layers_per_virtual_pipeline_stage is not None: + num_layers_per_pipeline_stage = args.num_layers_per_virtual_pipeline_stage + + if len(tokens_per_expert) != num_layers_per_pipeline_stage: + raise ValueError( + f'Expected {num_layers_per_pipeline_stage} token_per_experts ' + f'but found {len(tokens_per_expert)}.\nnum_layers = ' + f'{args.num_layers}\npipeline_model_parallel_size = ' + f'{args.pipeline_model_parallel_size}\n' + 'num_layers_per_virtual_pipeline_stage' + f' = {args.num_layers_per_virtual_pipeline_stage}', + ) + if len(expert_scores) != num_layers_per_pipeline_stage: + raise ValueError( + f'Expected {num_layers_per_pipeline_stage} expert_scores ' + f'but found {len(tokens_per_expert)}.\nnum_layers = ' + f'{args.num_layers}\npipeline_model_parallel_size = ' + f'{args.pipeline_model_parallel_size}\n' + 'num_layers_per_virtual_pipeline_stage' + f' = {args.num_layers_per_virtual_pipeline_stage}', + ) + + # Verify the shape of the tokens_per_expert and expert_scores tensors. + assert all((x.ndim == 1 and x.numel() == args.moe_num_experts for x in tokens_per_expert)) + + tokens = expert_scores[0].shape[0] + assert all(((x.ndim == 2 and x.shape[1] == args.moe_num_experts and x.shape[0] == tokens) for x in expert_scores)) + + # Concatenate the contributions of each layer and convert to + # the correct types and formats for the dot product. + expert_scores = torch.cat(expert_scores, dim=1) + if args.moe_lbl_in_fp32: + expert_scores = expert_scores.float() + if tokens != 0: + expert_scores = expert_scores.mean(dim=0) + else: + expert_scores = expert_scores.sum(dim=0) + tokens_per_expert = torch.cat(tokens_per_expert).to(expert_scores.dtype) + + expected_values = num_layers_per_pipeline_stage * args.moe_num_experts + assert tokens_per_expert.numel() == expected_values + assert expert_scores.numel() == expected_values + + # Calculate the total scale across all factors. + # + # loss_weight * num_experts / (num_layers * tokens * top_k) + scale_numerator = (args.moe_num_experts * args.moe_loss_weight) + scale_denominator = (args.num_layers * tokens * args.moe_top_k) + scale = scale_numerator / scale_denominator + return scale * torch.dot(tokens_per_expert, expert_scores) + + +# NOTE: This class defines MoE expert computation, including expert model parallel +# communication. When using FSDP on top of MegaBlocks this is the module that should +# be wrapped s.t. the weight all-gathers can be scheduled *before* the expert model +# parallel all2all. +class ParallelMLP(torch.nn.Module): + + def __init__(self, args: Arguments): + super(ParallelMLP, self).__init__() + self.args = args + + # Calculate the number of experts in total and the number of experts + # owned by this rank. + # world_size = mpu.get_expert_parallel_world_size(args) + self.num_experts = args.moe_num_experts + self.top_k = self.args.moe_top_k + + # Calculate the number of bits needed to represent the expert indices + # so that we can pass it to radix sort. + self.sort_end_bit = max(int(np.ceil(np.log2(self.num_experts))), 1) + + # Expert MLP. + self.mlp = mlp.MLP(args) + + self.bias: Optional[torch.Tensor] + if self.args.bias: + # Note that the output bias is not parallelized with expert + # model parallelism. + self.bias = torch.nn.Parameter( + torch.empty( + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + torch.nn.init.zeros_(self.bias) + else: + self.register_parameter('bias', None) + + # Select the forward function for the operating mode. + self.forward_fn = (self.parallel_forward_once if args.moe_expert_model_parallelism else self.forward_once) + + def expert_capacity(self, tokens: int) -> int: + world_size = mpu.get_expert_parallel_world_size(self.args) + tokens_per_expert = (self.top_k * tokens * world_size / self.num_experts) + return int(self.args.moe_capacity_factor * tokens_per_expert) + + def load_balancing_loss(self, tokens_per_expert: torch.Tensor, expert_scores: torch.Tensor): + """Calculate the load balancing loss contribution.""" + assert len(expert_scores.size()) == 2 + tokens, num_experts = expert_scores.size() + assert num_experts == self.num_experts + assert len(tokens_per_expert.size()) == 1 + num_experts, = tokens_per_expert.size() + assert num_experts == self.num_experts + scale = self.num_experts / (tokens * self.top_k) + return scale * torch.dot( + tokens_per_expert.to(expert_scores.dtype), + expert_scores.mean(dim=0), + ) + + def indices_and_bins(self, + top_expert: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # Sort the expert ids to produce the scatter/gather + # indices for the permutation. + # + # TODO(tgale): Is it worth doing this conversion to 32-bit + # prior? Could we place the `torch.max` operation to return + # 32-bit expert indices? + top_expert = top_expert.int() + # output = ops.sort(top_expert, self.sort_end_bit) + output = sort(top_expert, self.sort_end_bit) + assert output is not None + bin_ids, indices = output + + # Histogram the expert ids to identify the number of + # tokens routed to each expert. + # + # TODO(tgale): Does the sorted data produce a more favorable + # data distribution for histogram? Or is the op parallelism + # worth more? + # tokens_per_expert = ops.histogram(top_expert, self.num_experts) + tokens_per_expert = histogram(top_expert, self.num_experts) + + # Calculate the bin bounds for the sorted tokens. + # bins = ops.inclusive_cumsum(tokens_per_expert, 0) + bins = inclusive_cumsum(tokens_per_expert, 0) + assert bins is not None + bins = bins.view(1) if not len(bins.size()) else bins + + assert isinstance(indices, torch.Tensor) + assert isinstance(bin_ids, torch.Tensor) + assert isinstance(bins, torch.Tensor) + assert isinstance(tokens_per_expert, torch.Tensor) + + return indices, bin_ids, bins, tokens_per_expert + + def permute_and_compute( + self, + x: torch.Tensor, + tokens_per_expert: int, # unused + indices: torch.Tensor, + bin_ids: torch.Tensor, # unused + expert_weights: torch.Tensor, + bins: torch.Tensor, + expert_capacity: int, + top_k: int, + ): + # Route the tokens for MoE computation. + x = x.view(-1, x.shape[-1]) + # output = ops.binned_gather(x, indices, bins, expert_capacity, top_k) + output = binned_gather(x, indices, bins, expert_capacity, top_k) + assert output is not None + x = output + + # Perform the expert computation. Note that we don't + # use biases for these linear operations. + x = self.mlp(x) + + # Un-route the data for the MoE output. + # return ops.binned_scatter(x, indices, expert_weights, bins, top_k) + return binned_scatter(x, indices, expert_weights, bins, top_k) + + + def forward_once(self, x: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor): + # x: [sl, bs, hs] + # expert_weights: [sl * bs, top-k] + # top_experts: [sl * bs, top-k] + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + with torch.no_grad(): + indices, bin_ids, bins, tokens_per_expert = (self.indices_and_bins(top_experts)) + + # If expert_capacity is set to zero, set the number of tokens + # per expert to the maximum we need to avoid dropping tokens. + sl, bs, _ = x.size() + expert_capacity = self.expert_capacity(sl * bs) + if expert_capacity == 0: + expert_capacity = torch.max(tokens_per_expert).item() + + x = self.permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capacity, + self.top_k, + ) + return x, tokens_per_expert + + def parallel_forward_once(self, x: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor): + # NOTE: This function implements the same computation as forward_once + # but with expert model parallelism. + # + # 1. Permute the tokens locally so that they are grouped by their + # expert assignments. This allows us to transfer all of the tokens + # for a remote device in one communication primitive. + # + # 2. Permute the tokens across the expert parallel devices. After + # this is completed each device has all of the tokens assigned to + # its set of experts in its local HBM. + # + # 3. Permute the tokens locally so that they are grouped by their + # expert assignement. After the distributed permutation the tokens + # are grouped by which device they came from. We re-order them + # locally to allow for efficient computation. + # + # After this series of permutations we compute the linear layers + # and then repeat these three steps in reverse to produce the final + # output. + # + # Compute the mapping of local tokens to experts. + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + with torch.no_grad(): + indices, bin_ids, bins, tokens_per_expert = (self.indices_and_bins(top_experts)) + + # If we're sharding the experts along the hidden dimension + # multiple devices own parts of the same sets of experts. + # Replicate the token counts so every device gets the counts. + # repeated_tokens_per_expert = ops.repeat( + repeated_tokens_per_expert = repeat( + tokens_per_expert, + (mpu.hidden_sharding_degree(self.args),), + ) + + # Pass token count information to the device on which the + # target expert resides. + parallel_tokens_per_expert = torch.empty_like(repeated_tokens_per_expert,) + tpe_handle = dist.all_to_all_single( + parallel_tokens_per_expert, + repeated_tokens_per_expert, + group=self.args.expert_parallel_group, + async_op=True, + ) + + # Permute locally and without any padding so that tokens for each + # parallel device are stored contiguously. + # + # This view updates the shape of the tensor from [sl, bs, hs] to + # [sl * bs, hs] prior to the permutation. + x = x.view(-1, x.shape[-1]) + # output = ops.gather(x, indices, bin_ids, bins, self.top_k) + output = gather(x, indices, bin_ids, bins, self.top_k) + assert output is not None + x = output + + # Compute the number of tokens that will be received from each + # device and permute the input data across the devices. + with torch.no_grad(): + tpe_handle.wait() + experts_per_rank = mpu.experts_per_rank(self.args) + + # Reshape to [world_size, num_experts_per_rank]. + world_size = mpu.get_expert_parallel_world_size(self.args) + repeated_tokens_per_expert = (repeated_tokens_per_expert.view(world_size, experts_per_rank)) + parallel_tokens_per_expert = (parallel_tokens_per_expert.view(world_size, experts_per_rank)) + + # TODO(tgale): It might be faster to do this on the GPU and + # then communicate the results back to the host. + send_counts = repeated_tokens_per_expert.cpu().sum(dim=-1) + parallel_tokens_per_expert_cpu = parallel_tokens_per_expert.cpu() + recv_counts = parallel_tokens_per_expert_cpu.sum(dim=-1) + + # Convert the send/recv counts to lists. + send_counts = send_counts.tolist() + recv_counts = recv_counts.tolist() + tokens_received = sum(recv_counts) + + # If we're sharding the experts along the hidden dimension + # multiple devices own parts of the same sets of experts. + # Replicate the token counts so devices that share experts + # get all of the tokens assigned to them. + # + # TODO(tgale): Fuse this into the prior, local permutation. + # x = ops.repeat(x, (mpu.hidden_sharding_degree(self.args), 1)) + x = repeat(x, (mpu.hidden_sharding_degree(self.args), 1)) + + # Start the cross-device permutation asynchronously so we can + # overlap communication with computation. + parallel_x, parallel_x_handle = all_to_all( + x, + recv_counts, + send_counts, + self.args.expert_parallel_group, + async_op=True, + ) + + with torch.no_grad(): + # After we do the cross-device permutation we have the tokens on the + # correct device but not yet grouped by expert because we received + # tokens from each device as contiguous chunks. To group the tokens + # for expert computation we'll do one more local permutation. The + # rest of this torch.no_grad() scope sets up the indices and bins + # for this permutation. + # replicate_bins = ops.inclusive_cumsum( + replicate_bins = inclusive_cumsum( + parallel_tokens_per_expert.flatten(), + 0, + ) + replicate_bins = (replicate_bins.view(1) if not len(replicate_bins.size()) else replicate_bins) + + # Construct the expert indices for the permuted tokens. + parallel_top_expert = torch.remainder( + torch.arange( + self.num_experts * mpu.hidden_sharding_degree(self.args), + dtype=torch.int32, + device=indices.device, + ), + mpu.experts_per_rank(self.args), + ) + # parallel_top_expert = ops.replicate( + parallel_top_expert = replicate( + parallel_top_expert.unsqueeze(dim=0), + replicate_bins, + tokens_received, + ).flatten() + + # TODO(tgale): The sort_end_bit here can be reduced. + # parallel_bin_ids, parallel_indices = ops.sort( + parallel_bin_ids, parallel_indices = sort( + parallel_top_expert, + self.sort_end_bit, + ) + + # Calculate the bins boundaries from the token counts. + parallel_tokens_per_expert = parallel_tokens_per_expert.sum( + dim=0, + dtype=torch.int, + ) + # parallel_bins = ops.inclusive_cumsum(parallel_tokens_per_expert, 0) + parallel_bins = inclusive_cumsum(parallel_tokens_per_expert, 0) + parallel_bins = (parallel_bins.view(1) if not len(parallel_bins.size()) else parallel_bins) + + # If expert_capacity is set to zero, set the number of tokens + # per expert to the maximum we need to avoid dropping tokens. + tokens, _ = x.size() + expert_capacity = self.expert_capacity(tokens) + if expert_capacity == 0: + expert_capacity = torch.max(parallel_tokens_per_expert).item() + + # Locally permute the tokens and perform the expert computation. + # Block to make sure that the cross-device permutation is complete. + if self.args.mlp_impl == 'grouped': + # GroupedMLP requires counts on CPU. We can use the tensor already + # moved to CPU for the prior all_to_all, which avoids an extra + # device synchronization. + parallel_tokens_per_expert = parallel_tokens_per_expert_cpu.sum( + dim=0, + dtype=torch.int, + ) + parallel_x_handle.wait() + parallel_x = self.permute_and_compute( + parallel_x, + parallel_tokens_per_expert, + parallel_indices, + parallel_bin_ids, + None, # expert_weights + parallel_bins, + expert_capacity, + top_k=1, + ) + + # Un-permute the tokens across the devices. + x, _ = all_to_all( + parallel_x, + send_counts, + recv_counts, + self.args.expert_parallel_group, + ) + + # Reduce along the hidden sharding to get the final outputs. + # + # TODO(tgale): Fuse this into the following local permutation. + shape = ( + mpu.hidden_sharding_degree(self.args), + -1, + self.args.hidden_size, + ) + # x = ops.sum(x.view(shape), dim=0) + x = x.view(shape).sum(dim=0) + + # Un-permute locally to setup for the next series of operations. + # x = ops.scatter(x, indices, bin_ids, expert_weights, bins, self.top_k) + x = scatter(x, indices, bin_ids, expert_weights, bins, self.top_k) + return x, tokens_per_expert.flatten() + + def forward(self, x: torch.Tensor, scores: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor): + in_shape = x.size() + + # Compute the experts. + x, tokens_per_expert = self.forward_fn(x, expert_weights, top_experts) + if self.training and self.args.moe_loss_weight > 0: + save_load_balancing_loss((tokens_per_expert, scores)) + x = x.view(in_shape) + if self.bias is not None: + if self.args.return_bias: + return x, self.bias + return x + self.bias + return x + + +class MoE(torch.nn.Module): + + def __init__(self, args: Arguments): + super(MoE, self).__init__() + + # Token router. + self.router = router.LearnedRouter(args) + + # Expert computation helper. + self.experts = self._init_experts_mlp(args) + + self.shared_expert = None + if args.shared_expert: + # SharedExpert computation helper. + self.shared_expert = sharedexpert_registry.get(args) + + def _init_experts_mlp(self, args: Arguments): + return ParallelMLP(args) + + def forward(self, x: torch.Tensor): + # NOTE: If we're going to cast the activations to lower precision + # do it before we permute the tokens to save bandwidth. + x = common.cast_if_autocast_enabled(x) + + # Compute the expert scores and assignments. + scores, expert_weights, top_experts = self.router(x) + + # Compute the experts. + out = self.experts(x, scores, expert_weights, top_experts) + if self.shared_expert is not None: + shared_expert_out = self.shared_expert(x) + out = self.shared_expert.add_experts_sharedexpert( + shared_expert_out, + out, + ) + return out diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/mpu.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/mpu.py new file mode 100644 index 0000000000000000000000000000000000000000..434e143ab42bf3f83406d69e9dd1f72777716e22 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/mpu.py @@ -0,0 +1,94 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Optional + +import torch +import torch.distributed as dist + +# from megablocks.layers.arguments import Arguments +from .arguments import Arguments + + +class MoeParam(torch.Tensor): + + def __init__(self): + super().__init__(self) + self.expert_model_parallel: bool + + +def is_moe_param(tensor: torch.Tensor) -> bool: + return hasattr(tensor, 'expert_model_parallel') + + +def get_expert_parallel_world_size(args: Arguments) -> int: + return (dist.get_world_size(args.expert_parallel_group) if args.moe_expert_model_parallelism else 1) + + +def get_expert_parallel_rank(args: Arguments) -> int: + return (dist.get_rank(args.expert_parallel_group) if args.moe_expert_model_parallelism else 0) + + +def set_expert_model_parallel_attributes( + tensor: torch.Tensor, + is_parallel: bool, +): + assert not hasattr(tensor, 'expert_model_parallel') + setattr(tensor, 'expert_model_parallel', is_parallel) + + +def param_is_expert_model_parallel(param: MoeParam) -> bool: + return (hasattr(param, 'expert_model_parallel') and param.expert_model_parallel) + + +def copy_expert_model_parallel_attributes( + destination_tensor: torch.Tensor, + source_tensor: torch.Tensor, +): + if hasattr(source_tensor, 'expert_model_parallel'): + setattr( + destination_tensor, + 'expert_model_parallel', + getattr(source_tensor, 'expert_model_parallel'), + ) + + +def synchronized_print(group: Optional[dist.ProcessGroup], *x: torch.Tensor): + world_size = dist.get_world_size(group) + rank = dist.get_rank(group) + for i in range(world_size): + dist.barrier(group) + if i == rank: + print(f'rank = {rank}', *x) + + +# Helpers for expert/tensor sharding. +def expert_sharding_degree(args: Arguments) -> int: + world_size = get_expert_parallel_world_size(args) + esd = min(world_size, args.moe_num_experts) + + if (args.moe_num_experts % esd) != 0: + raise ValueError(f'Cannot shard {args.moe_num_experts} experts {esd} ways.',) + return esd + + +def hidden_sharding_degree(args: Arguments) -> int: + world_size = get_expert_parallel_world_size(args) + esd = expert_sharding_degree(args) + hsd = world_size // esd + + if (args.ffn_hidden_size % hsd) != 0: + raise ValueError(f'Cannot shard {args.ffn_hidden_size} features {hsd} ways.',) + if (esd * hsd) != world_size: + raise ValueError( + f"Invalid sharding. 'expert_sharding_degree' ({esd}) * hidden_sharding_degree ({hsd}) != world_size ({world_size}).", + ) + return hsd + + +def experts_per_rank(args: Arguments) -> int: + return args.moe_num_experts // expert_sharding_degree(args) + + +def features_per_rank(args: Arguments) -> int: + return args.ffn_hidden_size // hidden_sharding_degree(args) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/router.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/router.py new file mode 100644 index 0000000000000000000000000000000000000000..37cb2782348d62583376f1a183c7ede83601216d --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/router.py @@ -0,0 +1,116 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch + +# from megablocks.layers import common +# from megablocks.layers.arguments import Arguments +from . import common +from .arguments import Arguments + +_ROUTER_LOGITS = [] + + +def _save_router_logits(logits: torch.Tensor, args: Arguments): + if args.moe_zloss_weight == 0: + return + global _ROUTER_LOGITS + _ROUTER_LOGITS.append(logits) + + +def clear_router_zloss(): + global _ROUTER_LOGITS + _ROUTER_LOGITS.clear() + + +def batched_router_zloss(args: Arguments): + global _ROUTER_LOGITS + + if args.moe_zloss_weight == 0: + import warnings + warnings.warn('Call to batched_router_zloss, but moe_zloss_weight=0') + return 0 + + logits_per_router = _ROUTER_LOGITS + + if args.moe_zloss_in_fp32: + logits_per_router = [logits.float() for logits in logits_per_router] + + unscaled_zloss_per_router = torch.stack([ + torch.logsumexp(logits, dim=1).square().mean() for logits in logits_per_router + ]) + + return args.moe_zloss_weight * unscaled_zloss_per_router + + +# NOTE: To enable end-to-end benchmarking without convergence we +# support a flag to force the router to assign tokens uniformly +# across the experts. We do this with a custom autograd operation +# so that PyTorch still executes the full set of router operation. +class _UniformExpertAssignment(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, num_experts: int): + out = torch.arange(x.numel(), dtype=x.dtype, device=x.device) + out = torch.remainder(out, num_experts) + return out.view(x.shape) + + +_uniform_expert_assignment = _UniformExpertAssignment.apply + + +class LearnedRouter(torch.nn.Module): + + def __init__(self, args: Arguments): + super().__init__() + self.args = args + + # Learned router parameters. + # + # NOTE: This weight matrix is not parallelized with expert model + # parallelism. Each device needs the entire router weight matrix + # so that it can route its batch of data correctly. + self.layer = torch.nn.Linear( + args.hidden_size, + args.moe_num_experts, + bias=False, + dtype=common.dtype(args), + device=args.device, + ) + args.init_method(self.layer.weight) + + def jitter(self, x: torch.Tensor): + low: float = 1.0 - self.args.moe_jitter_eps + high: float = 1.0 + self.args.moe_jitter_eps + noise = torch.rand(x.size(), dtype=x.dtype, device=x.device) + return low + noise * (high - low) + + def _top_k(self, scores: torch.Tensor): + if self.args.moe_top_k == 1: + return scores.max(dim=-1, keepdim=True) + return torch.topk(scores, self.args.moe_top_k, dim=-1) + + def forward(self, x: torch.Tensor): + if self.training and self.args.moe_jitter_eps is not None: + x = x * self.jitter(x) + + logits = self.layer(x.view(-1, x.shape[-1])) + _save_router_logits(logits, self.args) + scores = logits.softmax(dim=-1) + expert_weights, expert_indices = self._top_k(scores) + if self.args.moe_normalize_expert_weights: + expert_weights = expert_weights / torch.norm( + expert_weights, + p=self.args.moe_normalize_expert_weights, + dim=-1, + keepdim=True, + ) + + expert_indices = ( + _uniform_expert_assignment( + expert_indices, + self.args.moe_num_experts, + ) if self.args.uniform_expert_assignment else expert_indices + ) + return scores, expert_weights, expert_indices diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_layers/sharedexpert_registry.py b/build/torch214-cxx11-cu126-x86_64-linux/_layers/sharedexpert_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..5840862f88f370ace5fd49bd0612fc98d186cc49 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_layers/sharedexpert_registry.py @@ -0,0 +1,32 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Union + +# from megablocks.layers import glu, mlp +# from megablocks.layers.arguments import Arguments +from . import glu, mlp +from .arguments import Arguments + +_REGISTRY = { + 'mlp': mlp.SharedMLP, + 'glu': glu.SharedGLU, +} + + +def get(args: Arguments) -> Union[mlp.SharedMLP, glu.SharedGLU]: + """Returns an SharedMLP for use in a dMoE instance. + + Uses the provided arguments to instantiate the appropriate + SharedMLP instance. + + Args: + args: propagated Arguments dataclass. + + Returns: + An instantiated SharedMLP constructed using the input args. + """ + if args.mlp_type not in _REGISTRY: + raise ValueError(f'Unsupported mlp type: {args.mlp_type}') + + return _REGISTRY[args.mlp_type](args) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_megablocks_cuda_addf474.abi3.so b/build/torch214-cxx11-cu126-x86_64-linux/_megablocks_cuda_addf474.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..c9257a4035f7fa64ee4c3ed6a67c0bd39d9820f4 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_megablocks_cuda_addf474.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:56f2ea87baee019a4710a34124e946d15f8823928e86d0366b480780e7e8542a +size 13822376 diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_ops.py b/build/torch214-cxx11-cu126-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..6640ac69d26eee39b52fe070e3c8571ed6a93d4e --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _megablocks_cuda_addf474 +ops = torch.ops._megablocks_cuda_addf474 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_megablocks_cuda_addf474::{op_name}" diff --git a/build/torch214-cxx11-cu126-x86_64-linux/_version.py b/build/torch214-cxx11-cu126-x86_64-linux/_version.py new file mode 100644 index 0000000000000000000000000000000000000000..c55783177af19bc03654c730c4892df8f8532279 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/_version.py @@ -0,0 +1,6 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +"""The MegaBlocks Version.""" + +__version__ = '0.11.0.dev0' diff --git a/build/torch214-cxx11-cu126-x86_64-linux/backend/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/backend/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9d4e43e9b7e6a9a1c3bde2df34914643ca5d8332 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/backend/__init__.py @@ -0,0 +1,2 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 diff --git a/build/torch214-cxx11-cu126-x86_64-linux/backend/kernels.py b/build/torch214-cxx11-cu126-x86_64-linux/backend/kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..d20139324883338ddb312e4b05a72056d74491ac --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/backend/kernels.py @@ -0,0 +1,557 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch +import triton +import triton.language as tl + +# Stub triton autotune when testing in a env that does not have CUDA +# this approach preserves the original code but enables testing without a GPU +if torch.cuda.is_available() is False: + import warnings + + warnings.warn("CUDA is not available. Triton autotuning is disabled.") + + def _no_autotune(*args, **kwargs): + def deco(fn): + return fn + return deco + + triton.autotune = _no_autotune + + +def assert_is_tensor(x, ndim): + if x.ndim != ndim: + raise ValueError(f'Expected {ndim}-tensor but got {x.ndim}-tensor') + + +def assert_is_matrix(x): + assert_is_tensor(x, 2) + + +def assert_is_vector(x): + if x.ndim != 1: + raise ValueError(f'Expected 1-tensor but got {x.ndim}-tensor') + + +def assert_equal(a, b): + if a != b: + raise ValueError(f'Expected dimensions to be equal but got {a} and {b}.',) + + +# a: (tokens, hidden_size), real. +# indices: (tokens * top_k), integer. +# bin_ids: (tokens * top_k), integer. +# weights: (tokens * top_k), real. +# bins: (num_experts), integer. +# padded_bins: (num_experts), integer. +@triton.autotune( + configs=[ + triton.Config({'BLOCK_X': 64}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=2), + triton.Config({'BLOCK_X': 256}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=4), + triton.Config({'BLOCK_X': 256}, num_warps=4), + ], + key=['NUM_COLUMNS'], +) +@triton.jit +def _padded_copy( + a, + b, + indices, + bin_ids, + weights, + bins, + padded_bins, + NUM_COLUMNS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_X: tl.constexpr, + A_TO_B: tl.constexpr, + SCALE: tl.constexpr, +): + # Our index into array 'a'. + index_a = tl.load(indices + tl.program_id(0)) + + # One threadblock per row in 'a'. Array 'b' has greater or equal + # number of rows since they could be padded. + bin_idx = tl.load(bin_ids + tl.program_id(0)) + + # Now we know what bin we're assigned to, but we need to know how + # many threadblocks were assigned to earlier bins so we can offset + # in our bin properly. + offset_in_bin = tl.program_id(0) + if bin_idx > 0: + offset_in_bin -= tl.load(bins + bin_idx - 1) + + # Load the starting index of our bin in array 'b'. + index_b = offset_in_bin + if bin_idx > 0: + index_b += tl.load(padded_bins + bin_idx - 1) + + # Offset the input and output pointers. + # + # If we're going from A to B, divide the input index to copy + # the same input repeatedly. If we're going from B to A we + # need to reduce the result. Using atomics is slow, so we + # do the reduce step in a second kernel. + offset = index_a // TOP_K if A_TO_B else index_a + a += tl.multiple_of(offset * NUM_COLUMNS, NUM_COLUMNS) + b += tl.multiple_of(index_b * NUM_COLUMNS, NUM_COLUMNS) + offsets = tl.max_contiguous(tl.arange(0, BLOCK_X), BLOCK_X) + + # Load the scale, if requested. + scale = tl.load(weights + index_a) if SCALE else 1 + + # Swap the pointers depending on the direction. + iptr = a if A_TO_B else b + optr = b if A_TO_B else a + + iterations = tl.cdiv(NUM_COLUMNS, BLOCK_X) + for _ in range(iterations): + mask = offsets < NUM_COLUMNS + x = tl.load(iptr + offsets, mask=mask) + x = x.to(tl.float32) * scale.to(tl.float32) + + tl.store(optr + offsets, x.to(optr.dtype.element_ty), mask=mask) + + offsets += BLOCK_X + + +def padded_gather(x, indices, bin_ids, weights, bins, padded_bins, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_vector(indices) + assert_is_vector(bin_ids) + assert_is_vector(bins) + assert_is_vector(padded_bins) + assert_equal(indices.shape[0], x.shape[0] * top_k) + assert_equal(bin_ids.shape[0], x.shape[0] * top_k) + assert_equal(bins.size(), padded_bins.size()) + + if weights is not None: + assert_equal(weights.shape[0], x.shape[0] * top_k) + + # NOTE: Because of the padding, the output size is dynamic. + # We load the final padded bin bound to get the output rows. + output_rows = padded_bins[-1].cpu().item() + out = torch.zeros((output_rows, x.shape[1]), dtype=x.dtype, device=x.device) + _padded_copy[(indices.shape[0],)]( + x, + out, + indices, + bin_ids, + weights, + bins, + padded_bins, + NUM_COLUMNS=x.shape[1], + A_TO_B=True, + TOP_K=top_k, + SCALE=weights is not None, + ) + return out + + +def gather(x, indices, bin_ids, weights, bins, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_vector(indices) + assert_is_vector(bin_ids) + assert_is_vector(bins) + assert_equal(indices.shape[0], x.shape[0] * top_k) + assert_equal(bin_ids.shape[0], x.shape[0] * top_k) + + if weights is not None: + assert_equal(weights.shape[0], x.shape[0] * top_k) + + # NOTE: There is no padding so the output rows equals the + # input rows multiplied by top_k. + output_rows = x.shape[0] * top_k + out = torch.empty((output_rows, x.shape[1]), dtype=x.dtype, device=x.device) + _padded_copy[(indices.shape[0],)]( + x, + out, + indices, + bin_ids, + weights, + bins, + bins, + NUM_COLUMNS=x.shape[1], + A_TO_B=True, + TOP_K=top_k, + SCALE=weights is not None, + ) + return out + + +def padded_scatter(x, indices, bin_ids, weights, bins, padded_bins, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_vector(indices) + assert_is_vector(bin_ids) + assert_is_vector(bins) + assert_is_vector(padded_bins) + assert_equal(indices.shape[0], bin_ids.shape[0]) + assert_equal(bins.size(), padded_bins.size()) + + if weights is not None: + assert_equal(indices.shape[0], weights.shape[0]) + + tokens = indices.shape[0] // top_k + out = torch.empty((tokens, top_k, x.shape[1]), dtype=x.dtype, device=x.device) + _padded_copy[(indices.shape[0],)]( + out, + x, + indices, + bin_ids, + weights, + bins, + padded_bins, + NUM_COLUMNS=x.shape[1], + A_TO_B=False, + TOP_K=top_k, + SCALE=weights is not None, + ) + + # Reduce along the top-k dimension, if needed. + return out.sum(dim=1) if top_k > 1 else out.view(tokens, x.shape[1]) + + +def scatter(x, indices, bin_ids, weights, bins, top_k): + return padded_scatter(x, indices, bin_ids, weights, bins, bins, top_k) + + +# x: (tokens, top_k, hidden_size), real +# grad: (tokens, hidden_size), real. +# wgrad: (tokens, top_k), real. +# indices: (tokens * top_k), integer. +# bin_ids: (tokens * top_k), integer. +# bins: (num_experts), integer. +# padded_bins: (num_experts), integer. +@triton.autotune( + configs=[ + triton.Config({'BLOCK_X': 64}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=2), + triton.Config({'BLOCK_X': 256}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=4), + triton.Config({'BLOCK_X': 256}, num_warps=4), + ], + key=['NUM_COLUMNS'], +) +@triton.jit +def _padded_copy_wgrad( + x, + grad, + wgrad, + indices, + bin_ids, + bins, + padded_bins, + NUM_COLUMNS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_X: tl.constexpr, +): + # Our index into 'tokens * top_k'. + index_out = tl.load(indices + tl.program_id(0)) + + # One threadblock per row in 'a'. Array 'b' has greater or equal + # number of rows since they could be padded. + bin_idx = tl.load(bin_ids + tl.program_id(0)) + + # Now we know what bin we're assigned to, but we need to know how + # many threadblocks were assigned to earlier bins so we can offset + # in our bin properly. + offset_in_bin = tl.program_id(0) + if bin_idx > 0: + offset_in_bin -= tl.load(bins + bin_idx - 1) + + # Load the starting index of our bin in array 'x'. + index_x = offset_in_bin + if bin_idx > 0: + index_x += tl.load(padded_bins + bin_idx - 1) + + # Offset the input and output pointers. + wgrad += index_out + grad += tl.multiple_of((index_out // TOP_K) * NUM_COLUMNS, NUM_COLUMNS) + x += tl.multiple_of(index_x * NUM_COLUMNS, NUM_COLUMNS) + offsets = tl.max_contiguous(tl.arange(0, BLOCK_X), BLOCK_X) + + acc = tl.zeros((BLOCK_X,), dtype=tl.float32) + iterations = tl.cdiv(NUM_COLUMNS, BLOCK_X) + for _ in range(iterations): + mask = offsets < NUM_COLUMNS + data = tl.load(x + offsets, mask=mask).to(tl.float32) + scale = tl.load(grad + offsets, mask=mask).to(tl.float32) + acc += data * scale + offsets += BLOCK_X + + # Reduce to get the final result and store. + out = tl.sum(acc).to(wgrad.dtype.element_ty) + tl.store(wgrad, out) + + +def padded_scatter_wgrad(x, grad, indices, bin_ids, bins, padded_bins, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_matrix(grad) + assert_is_vector(indices) + assert_is_vector(bin_ids) + assert_is_vector(bins) + assert_is_vector(padded_bins) + assert_equal(indices.shape[0], bin_ids.shape[0]) + assert_equal(bins.size(), padded_bins.size()) + + tokens = indices.shape[0] // top_k + out = torch.empty((tokens * top_k), dtype=x.dtype, device=x.device) + _padded_copy_wgrad[(indices.shape[0],)]( + x, + grad, + out, + indices, + bin_ids, + bins, + padded_bins, + NUM_COLUMNS=x.shape[1], + TOP_K=top_k, + ) + return out + + +def scatter_wgrad(x, grad, indices, bin_ids, bins, top_k): + return padded_scatter_wgrad(x, grad, indices, bin_ids, bins, bins, top_k) + + +# a: (tokens, hidden_size), real. +# b: (num_experts, expert_capacity, num_columns), real. +# indices: (tokens * top_k), integer. +# weights: (tokens * top_k), real. +# bins: (num_experts), integer. +@triton.autotune( + configs=[ + triton.Config({'BLOCK_X': 64}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=2), + triton.Config({'BLOCK_X': 256}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=4), + triton.Config({'BLOCK_X': 256}, num_warps=4), + ], + key=['NUM_COLUMNS'], +) +@triton.jit +def _binned_copy( + a, + b, + num_experts, + expert_capacity, + indices, + weights, + bins, + NUM_COLUMNS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_X: tl.constexpr, + A_TO_B: tl.constexpr, + SCALE: tl.constexpr, +): + # Load our indices into the output. + expert_idx = tl.program_id(0) + entry_idx = tl.program_id(1) + + # Calculate our offset into the output. + index_b = expert_idx * expert_capacity + entry_idx + + # Load the index bounds for our bin and calculate + # the number of tokens assigned to our expert. + start = 0 + if expert_idx > 0: + start = tl.load(bins + expert_idx - 1) + end = tl.load(bins + expert_idx) + num_tokens = end - start + + # Calculate our offset into the input. If we don't + # have an input exit early. + if entry_idx >= num_tokens: + return + index_a = tl.load(indices + start + entry_idx) + + # Offset the input and output pointers. + # + # If we're going from A to B, divide the input index to copy + # the same input repeatedly. If we're going from B to A we + # need to reduce the result. Using atomics is slow, so we + # do the reduce step in a second kernel. + offset = index_a // TOP_K if A_TO_B else index_a + a += tl.multiple_of(offset * NUM_COLUMNS, NUM_COLUMNS) + b += tl.multiple_of(index_b * NUM_COLUMNS, NUM_COLUMNS) + offsets = tl.max_contiguous(tl.arange(0, BLOCK_X), BLOCK_X) + + # Load the scale, if requested. + scale = tl.load(weights + index_a) if SCALE else 1 + + # Swap the pointers depending on the direction. + # + # NOTE: We need to zero the output in both directions. + iptr = a if A_TO_B else b + optr = b if A_TO_B else a + + iterations = tl.cdiv(NUM_COLUMNS, BLOCK_X) + for _ in range(iterations): + mask = offsets < NUM_COLUMNS + x = tl.load(iptr + offsets, mask=mask) + x = x.to(tl.float32) * scale.to(tl.float32) + + tl.store(optr + offsets, x.to(optr.dtype.element_ty), mask=mask) + + offsets += BLOCK_X + + +def binned_gather(x, indices, weights, bins, expert_capacity, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_vector(indices) + assert_is_vector(bins) + assert_equal(indices.shape[0], x.shape[0] * top_k) + + if weights is not None: + assert_equal(weights.shape[0], x.shape[0] * top_k) + + num_experts = bins.shape[0] + out = torch.zeros((num_experts, expert_capacity, x.shape[1]), dtype=x.dtype, device=x.device) + + _binned_copy[(num_experts, expert_capacity)]( + x, + out, + num_experts, + expert_capacity, + indices, + weights, + bins, + NUM_COLUMNS=x.shape[1], + A_TO_B=True, + TOP_K=top_k, + SCALE=weights is not None, + ) + return out + + +def binned_scatter(x, indices, weights, bins, top_k): + # Validate the input shapes. + assert_is_tensor(x, 3) + assert_is_vector(indices) + assert_is_vector(bins) + assert_equal(bins.shape[0], x.shape[0]) + + if weights is not None: + assert_equal(indices.shape[0], weights.shape[0]) + + num_experts, expert_capacity, hidden_size = x.shape + tokens = indices.shape[0] // top_k + out = torch.zeros((tokens, top_k, hidden_size), dtype=x.dtype, device=x.device) + _binned_copy[(num_experts, expert_capacity)]( + out, + x, + num_experts, + expert_capacity, + indices, + weights, + bins, + NUM_COLUMNS=hidden_size, + A_TO_B=False, + TOP_K=top_k, + SCALE=weights is not None, + ) + + # Reduce along the top-k dimension, if needed. + return out.sum(dim=1) if top_k > 1 else out.view(tokens, hidden_size) + + +# a: (tokens, hidden_size), real. +# b: (num_experts, expert_capacity, num_columns), real. +# indices: (tokens * top_k), integer. +# weights: (tokens * top_k), real. +# bins: (num_experts), integer. +@triton.autotune( + configs=[ + triton.Config({'BLOCK_X': 64}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=2), + triton.Config({'BLOCK_X': 256}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=4), + triton.Config({'BLOCK_X': 256}, num_warps=4), + ], + key=['NUM_COLUMNS'], +) +@triton.jit +def _binned_copy_wgrad( + x, + grad, + wgrad, + num_experts, + expert_capacity, + indices, + bins, + NUM_COLUMNS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_X: tl.constexpr, +): + # Load our indices into the output. + expert_idx = tl.program_id(0) + entry_idx = tl.program_id(1) + + # Calculate our offset into the output. + index_x = expert_idx * expert_capacity + entry_idx + + # Load the index bounds for our bin and calculate + # the number of tokens assigned to our expert. + start = 0 + if expert_idx > 0: + start = tl.load(bins + expert_idx - 1) + end = tl.load(bins + expert_idx) + num_tokens = end - start + + # Calculate our offset into the input. If we don't + # have an input exit early. + if entry_idx >= num_tokens: + return + index_out = tl.load(indices + start + entry_idx) + + # Offset the input and output pointers. + wgrad += index_out + grad += tl.multiple_of((index_out // TOP_K) * NUM_COLUMNS, NUM_COLUMNS) + x += tl.multiple_of(index_x * NUM_COLUMNS, NUM_COLUMNS) + offsets = tl.max_contiguous(tl.arange(0, BLOCK_X), BLOCK_X) + + acc = tl.zeros((BLOCK_X,), dtype=tl.float32) + iterations = tl.cdiv(NUM_COLUMNS, BLOCK_X) + for _ in range(iterations): + mask = offsets < NUM_COLUMNS + data = tl.load(x + offsets, mask=mask).to(tl.float32) + scale = tl.load(grad + offsets, mask=mask).to(tl.float32) + acc += data * scale + offsets += BLOCK_X + + # Reduce to get the final result and store. + out = tl.sum(acc).to(wgrad.dtype.element_ty) + tl.store(wgrad, out) + + +def binned_scatter_wgrad(x, grad, indices, bins, top_k): + # Validate the input shapes. + assert_is_tensor(x, 3) + assert_is_matrix(grad) + assert_is_vector(indices) + assert_is_vector(bins) + assert_equal(bins.shape[0], x.shape[0]) + + num_experts, expert_capacity, hidden_size = x.shape + tokens = indices.shape[0] // top_k + out = torch.zeros((tokens * top_k), dtype=x.dtype, device=x.device) + _binned_copy_wgrad[(num_experts, expert_capacity)]( + x, + grad, + out, + num_experts, + expert_capacity, + indices, + bins, + NUM_COLUMNS=hidden_size, + TOP_K=top_k, + ) + return out diff --git a/build/torch214-cxx11-cu126-x86_64-linux/benchmark_util.py b/build/torch214-cxx11-cu126-x86_64-linux/benchmark_util.py new file mode 100644 index 0000000000000000000000000000000000000000..02612d95e3ead1175a596e2878fa34b5bf85ad6f --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/benchmark_util.py @@ -0,0 +1,35 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import torch + + +def log_benchmark(name, arguments, time, std): + print('=' * 60) + print(f'{name} Benchmark') + print('Benchmark Parameters:') + for (key, value) in arguments.items(): + print(f'{key} = {value}') + print('Results:') + print('mean time = {:.3f}ms, std time = {:.3f}ms'.format(time, std)) + print('=' * 60) + + +def benchmark_function(fn, iterations=100, warmup=10): + # Warmup iterations. + for _ in range(warmup): + fn() + + times = [] + for i in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + start.record() + fn() + end.record() + + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + return np.mean(times), np.std(times) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/cpu_fused_moe.py b/build/torch214-cxx11-cu126-x86_64-linux/cpu_fused_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..7a29e81d6e9c57f64b0a78ac1c0828e45fd9d855 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/cpu_fused_moe.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# MegaBlocks CPU Fused MoE Implementation +# +# This is a pure Python/PyTorch implementation for CPU. +# For better performance, consider using the C++ kernel implementation. +# +import torch +import torch.nn.functional as F + + +def swigluoai_activation(gate: torch.Tensor, up: torch.Tensor, + alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor: + """ + SwigluOAI activation function used in GptOss models. + + Formula: + gate = clamp(gate, max=limit) + up = clamp(up, -limit, limit) + glu = gate * sigmoid(gate * alpha) + output = (up + 1) * glu + + Args: + gate: Gate tensor from gate projection + up: Up tensor from up projection + alpha: Scaling factor for sigmoid (default: 1.702) + limit: Clamp limit (default: 7.0) + + Returns: + Activated tensor + """ + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + glu = gate * torch.sigmoid(gate * alpha) + return (up + 1) * glu + + +def silu_and_mul_activation(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """ + SiLU (Swish) activation with element-wise multiplication. + + Formula: + output = silu(gate) * up + + Args: + gate: Gate tensor + up: Up tensor + + Returns: + Activated tensor + """ + return F.silu(gate) * up + + +def route_tokens_cpu( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor | None, + moe_top_k: int, + moe_num_experts: int, + moe_normalize_expert_weights: int | None = None, +) -> tuple: + """ + Route tokens to experts and compute expert weights and indices (CPU version). + + Args: + x: Input tensor [batch, seq, hidden] or [tokens, hidden] + router_weight: Router weight [num_experts, hidden] + router_bias: Router bias [num_experts] or None + moe_top_k: Number of experts per token + moe_num_experts: Total number of experts + moe_normalize_expert_weights: Normalization order or None + + Returns: + Tuple of (logits, expert_weights, expert_indices) + """ + x_flat = x.view(-1, x.shape[-1]) + logits = F.linear(x_flat, router_weight, router_bias) + + if moe_top_k == 1: + expert_weights, expert_indices = logits.max(dim=-1, keepdim=True) + else: + expert_weights, expert_indices = torch.topk(logits, moe_top_k, dim=-1) + + expert_weights = expert_weights.softmax(dim=-1) + + if moe_normalize_expert_weights is not None: + expert_weights = expert_weights / torch.norm( + expert_weights, + p=moe_normalize_expert_weights, + dim=-1, + keepdim=True, + ) + + return logits, expert_weights, expert_indices + + +def cpu_fused_moe( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + activation: str = "silu", + alpha: float = 1.702, + limit: float = 7.0, + is_interleaved: bool = True, +) -> torch.Tensor: + """ + CPU Fused MoE using PyTorch operations. + + This implementation processes all experts in parallel using batched operations + instead of sequential for loops, which is more efficient on CPU. + + Args: + hidden_states: [num_tokens, hidden_size] + w1: [num_experts, hidden_size, 2*inter_size] - gate_up_proj weights + w2: [num_experts, inter_size, hidden_size] - down_proj weights + topk_weights: [num_tokens, topk] - routing weights + topk_ids: [num_tokens, topk] - expert indices + w1_bias: [num_experts, 2*inter_size] or None + w2_bias: [num_experts, hidden_size] or None + activation: "silu" or "swigluoai" + alpha: swigluoai alpha parameter + limit: swigluoai limit parameter + is_interleaved: whether gate_up is interleaved [g0,u0,g1,u1,...] (True for GptOss) + + Returns: + output: [num_tokens, hidden_size] + """ + num_tokens, hidden_size = hidden_states.shape + num_experts = w1.shape[0] + inter_size = w2.shape[1] + topk = topk_weights.shape[1] + + # Initialize output + output = torch.zeros_like(hidden_states) + + # Build expert mask: which tokens go to which expert + # expert_mask[expert_id] contains indices of (token_idx, topk_pos) pairs + for expert_idx in range(num_experts): + # Find tokens assigned to this expert + # mask shape: [num_tokens, topk], True where topk_ids == expert_idx + mask = (topk_ids == expert_idx) + + if not mask.any(): + continue + + # Get token indices and topk positions + token_indices, topk_positions = torch.where(mask) + + if len(token_indices) == 0: + continue + + # Gather input tokens for this expert + # current_hidden: [num_selected_tokens, hidden_size] + current_hidden = hidden_states[token_indices] + + # Get weights for this expert + # w1[expert_idx]: [hidden_size, 2*inter_size] + # w2[expert_idx]: [inter_size, hidden_size] + expert_w1 = w1[expert_idx] # [hidden_size, 2*inter_size] + expert_w2 = w2[expert_idx] # [inter_size, hidden_size] + + # First projection: hidden @ w1 -> [num_selected, 2*inter_size] + gate_up = current_hidden @ expert_w1 + + # Add bias if present + if w1_bias is not None: + gate_up = gate_up + w1_bias[expert_idx] + + # Split gate and up projections + if is_interleaved: + # GptOss uses interleaved layout: [g0, u0, g1, u1, ...] + gate = gate_up[..., ::2] # [num_selected, inter_size] + up = gate_up[..., 1::2] # [num_selected, inter_size] + else: + # Standard layout: [gate_all, up_all] + gate = gate_up[..., :inter_size] + up = gate_up[..., inter_size:] + + # Apply activation + if activation == "swigluoai": + activated = swigluoai_activation(gate, up, alpha, limit) + else: # silu + activated = silu_and_mul_activation(gate, up) + + # Second projection: activated @ w2 -> [num_selected, hidden_size] + expert_out = activated @ expert_w2 + + # Add bias if present + if w2_bias is not None: + expert_out = expert_out + w2_bias[expert_idx] + + # Apply routing weights and accumulate + # weights shape: [num_selected] + weights = topk_weights[token_indices, topk_positions].unsqueeze(-1) + weighted_out = expert_out * weights + + # Accumulate to output + output.index_add_(0, token_indices, weighted_out.to(output.dtype)) + + return output + + +class MegaBlocksMoeMLP(torch.nn.Module): + """ + CPU MoE MLP module that can be used as a drop-in replacement for + the transformers GptOssMLP when using @use_kernel_forward_from_hub. + """ + can_torch_compile: bool = True + + def forward(self, x: torch.Tensor) -> tuple: + """ + Forward pass through the MoE layer. + + Args: + x: Input tensor of shape [batch_size, seq_len, hidden_size] or [tokens, hidden_size] + + Returns: + Tuple of (output, expert_weights) where: + - output: Tensor of same shape as input + - expert_weights: Expert weights for each token [tokens, top_k] + """ + # Get MoE parameters from the wrapped modules + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + moe_normalize_expert_weights = getattr( + self.experts, "normalize_expert_weights", None + ) + + # Detect activation type + if hasattr(self.experts, "alpha") and hasattr(self.experts, "limit"): + activation = "swigluoai" + alpha = self.experts.alpha + limit = self.experts.limit + else: + activation = getattr(self.experts, "activation", "silu") + alpha = 1.702 + limit = 7.0 + + # Get weight tensors + if hasattr(self.experts, "gate_up_proj"): + w1 = self.experts.gate_up_proj + is_interleaved = True # GptOss uses interleaved layout + elif hasattr(self.experts, "w1"): + w1 = self.experts.w1 + w3 = getattr(self.experts, "w3", None) + if w3 is not None: + w1 = torch.cat([w1, w3], dim=-1) + is_interleaved = False + else: + raise AttributeError("experts module must have 'gate_up_proj' or 'w1' attribute") + + if hasattr(self.experts, "down_proj"): + w2 = self.experts.down_proj + elif hasattr(self.experts, "w2"): + w2 = self.experts.w2 + else: + raise AttributeError("experts module must have 'down_proj' or 'w2' attribute") + + # Get optional bias tensors + w1_bias = getattr(self.experts, "gate_up_proj_bias", None) + w2_bias = getattr(self.experts, "down_proj_bias", None) + + # Store original shape + in_shape = x.size() + + # Route tokens to experts + logits, expert_weights, expert_indices = route_tokens_cpu( + x, + self.router.weight, + getattr(self.router, "bias", None), + moe_top_k, + moe_num_experts, + moe_normalize_expert_weights, + ) + + # Reshape input for fused MoE + x_flat = x.view(-1, x.shape[-1]) + + # Call CPU fused MoE + output = cpu_fused_moe( + hidden_states=x_flat, + w1=w1, + w2=w2, + topk_weights=expert_weights, + topk_ids=expert_indices, + w1_bias=w1_bias, + w2_bias=w2_bias, + activation=activation, + alpha=alpha, + limit=limit, + is_interleaved=is_interleaved, + ) + + # Restore original shape + output = output.view(in_shape) + + return output, expert_weights + + +# Export classes and functions +__all__ = [ + "MegaBlocksMoeMLP", + "cpu_fused_moe", + "route_tokens_cpu", + "swigluoai_activation", + "silu_and_mul_activation", +] diff --git a/build/torch214-cxx11-cu126-x86_64-linux/cpu_moe_cpp.py b/build/torch214-cxx11-cu126-x86_64-linux/cpu_moe_cpp.py new file mode 100644 index 0000000000000000000000000000000000000000..073ff66d24ce348fbb5ed19c9027fadd3f7a9c61 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/cpu_moe_cpp.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +# MegaBlocks C++ Optimized CPU MoE + +""" +C++ accelerated MoE with brgemm optimization for Intel AMX. +Direct replacement for cpu_fused_moe.MegaBlocksMoeMLP with better performance. +""" + +import torch +from typing import Optional +from .cpu_fused_moe import route_tokens_cpu +from ._ops import ops + + +def _to_local_tensor(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + """Convert DTensor to local torch.Tensor if needed for custom ops compatibility.""" + if tensor is None: + return None + # Check if it's a DTensor by looking for the to_local() method + if hasattr(tensor, "to_local"): + return tensor.to_local() + return tensor + + +def fused_moe_cpp( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + inplace: bool = False, + use_int8_w8a8: bool = False, + use_fp8_w8a16: bool = False, + use_mxfp4: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + block_size: Optional[list] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + w1_bias: Optional[torch.Tensor] = None, + w2_bias: Optional[torch.Tensor] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + is_vnni: bool = False, +) -> torch.Tensor: + """ + C++ Fused MoE with brgemm optimization (sglang compatible interface). + + Uses at::native::cpublas::brgemm for efficient batch GEMM on Intel CPUs. + Supports both silu_and_mul (standard SwiGLU) and swigluoai (GptOss) activations. + + Args: + hidden_states: Input tensor [M, K] + w1: Gate and up projections [E, 2N, K] + w2: Down projection [E, K, N] + topk_weights: Expert weights [M, topk] + topk_ids: Expert indices [M, topk] + inplace: Whether to use hidden_states as output + use_int8_w8a8: Use int8 quantization + use_fp8_w8a16: Use fp8 quantization + use_mxfp4: Use mxfp4 quantization + w1_scale, w2_scale: Quantization scales + block_size: Block size for fp8 + a1_scale, a2_scale: Activation scales + w1_bias, w2_bias: Optional biases + alpha: swigluoai alpha parameter (set to enable swiglu) + limit: swigluoai limit parameter (set to enable swiglu) + is_vnni: Whether w1/w2 are already in VNNI packed format + """ + # MXFP4/FP8 kernels only support bf16, convert if needed + orig_dtype = hidden_states.dtype + need_convert = ((use_mxfp4 or use_fp8_w8a16) and orig_dtype != torch.bfloat16) or orig_dtype == torch.float32 + if need_convert: + hidden_states = hidden_states.to(torch.bfloat16) + + # bias must match hidden_states dtype + if w1_bias is not None: + w1_bias = w1_bias.to(hidden_states.dtype) + if w2_bias is not None: + w2_bias = w2_bias.to(hidden_states.dtype) + + # Convert DTensor to local tensor for custom ops compatibility (TP mode) + hidden_states = _to_local_tensor(hidden_states) + w1 = _to_local_tensor(w1) + w2 = _to_local_tensor(w2) + topk_weights = _to_local_tensor(topk_weights) + topk_ids = _to_local_tensor(topk_ids) + w1_scale = _to_local_tensor(w1_scale) + w2_scale = _to_local_tensor(w2_scale) + a1_scale = _to_local_tensor(a1_scale) + a2_scale = _to_local_tensor(a2_scale) + w1_bias = _to_local_tensor(w1_bias) + w2_bias = _to_local_tensor(w2_bias) + + output = ops.fused_experts( + hidden_states, w1, w2, topk_weights, topk_ids, + inplace, use_int8_w8a8, use_fp8_w8a16, use_mxfp4, + w1_scale, w2_scale, block_size, a1_scale, a2_scale, + w1_bias, w2_bias, alpha, limit, is_vnni + ) + + # Convert back to original dtype if needed + if need_convert: + output = output.to(orig_dtype) + return output + + +class CPUMegaBlocksMoeMLP(torch.nn.Module): + """ + C++ optimized MoE MLP using brgemm. + Drop-in replacement for cpu_fused_moe.MegaBlocksMoeMLP with better performance. + + Usage in transformers: + # Will be used via @use_kernel_forward_from_hub decorator + """ + can_torch_compile: bool = True + + def forward(self, x: torch.Tensor) -> tuple: + """ + Forward pass through the MoE layer using C++ kernel. + + Args: + x: Input tensor [batch_size, seq_len, hidden_size] + + Returns: + Tuple of (output, expert_weights) + """ + # Optimization for GPT-OSS model + if getattr(self, "use_mxfp4", None) is None: + self.use_mxfp4 = False + + w1_scale = None + w2_scale = None + + if ( + not getattr(self, "packed_scales", False) + and hasattr(self.experts, "gate_up_proj") + and getattr(self.experts, "gate_up_proj_precision_config", None) is not None + ): + # convert scales + data_1 = ops.convert_scale_packed(self.experts.gate_up_proj_precision_config.weight_scale.data.transpose(-1, -2).contiguous()) + data_2 = ops.convert_scale_packed(self.experts.down_proj_precision_config.weight_scale.data.transpose(-1, -2).contiguous()) + self.experts.gate_up_proj_precision_config.weight_scale.storage.data = data_1 + self.experts.down_proj_precision_config.weight_scale.storage.data = data_2 + self.packed_scales = True + self.use_mxfp4 = True + + if not getattr(self, "packed_weight", False) and hasattr( + self.experts, "gate_up_proj" + ): + # convert weights + data_1 = self.experts.gate_up_proj.data.transpose(-1, -2).contiguous() + data_2 = self.experts.down_proj.data.transpose(-1, -2).contiguous() + if self.use_mxfp4: + self.experts.gate_up_proj.storage.data = ops.convert_weight_packed(data_1) + self.experts.down_proj.storage.data = ops.convert_weight_packed(data_2) + else: + # convert_weight_packed only supports bfloat16, float16, int8, fp8_e4m3 or uint8(mxfp4 or int4). + data_1 = data_1.to(torch.bfloat16) if data_1.dtype == torch.float32 else data_1 + data_2 = data_2.to(torch.bfloat16) if data_2.dtype == torch.float32 else data_2 + self.experts.gate_up_proj.data = ops.convert_weight_packed(data_1) + self.experts.down_proj.data = ops.convert_weight_packed(data_2) + + # C++ kernel does not support float32. + dtype = torch.bfloat16 if x.dtype == torch.float32 else x.dtype + if getattr(self.experts, "gate_up_proj_bias", None) is not None: + self.experts.gate_up_proj_bias.data = self.experts.gate_up_proj_bias.data.to(dtype) + if getattr(self.experts, "down_proj_bias", None) is not None: + self.experts.down_proj_bias.data = self.experts.down_proj_bias.data.to(dtype) + + self.packed_weight = True + + # Get MoE parameters + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None) + + # Detect activation type + if hasattr(self.experts, "alpha") and hasattr(self.experts, "limit"): + activation = "swigluoai" + alpha = self.experts.alpha + limit = self.experts.limit + else: + activation = getattr(self.experts, "activation", "silu") + alpha = 1.702 + limit = 7.0 + + # Get weight tensors + if hasattr(self.experts, "gate_up_proj"): + w1 = self.experts.gate_up_proj + elif hasattr(self.experts, "w1"): + w1 = self.experts.w1 + w3 = getattr(self.experts, "w3", None) + if w3 is not None: + w1 = torch.cat([w1, w3], dim=-1) + else: + raise AttributeError("experts module must have 'gate_up_proj' or 'w1' attribute") + + if hasattr(self.experts, "down_proj"): + w2 = self.experts.down_proj + elif hasattr(self.experts, "w2"): + w2 = self.experts.w2 + else: + raise AttributeError("experts module must have 'down_proj' or 'w2' attribute") + + # Get optional bias tensors + w1_bias = getattr(self.experts, "gate_up_proj_bias", None) + w2_bias = getattr(self.experts, "down_proj_bias", None) + w1_bias = w1_bias if w1_bias is None else w1_bias.data + w2_bias = w2_bias if w2_bias is None else w2_bias.data + + if self.use_mxfp4: + w1_scale = self.experts.gate_up_proj_precision_config.weight_scale.data + w2_scale = self.experts.down_proj_precision_config.weight_scale.data + + # Store original shape + in_shape = x.size() + + # Route tokens to experts (Python implementation is fast enough) + logits, expert_weights, expert_indices = route_tokens_cpu( + x, + self.router.weight, + getattr(self.router, "bias", None), + moe_top_k, + moe_num_experts, + moe_normalize_expert_weights, + ) + + # Flatten input + x_flat = x.view(-1, x.shape[-1]) + + # Determine alpha/limit for swiglu activation + use_alpha = alpha if activation == "swigluoai" else None + use_limit = limit if activation == "swigluoai" else None + + # Call C++ optimized kernel + output = fused_moe_cpp( + hidden_states=x_flat, + w1=w1.data, + w2=w2.data, + topk_weights=expert_weights, + topk_ids=expert_indices.to(torch.int32), + inplace=False, + use_int8_w8a8=False, + use_fp8_w8a16=False, + use_mxfp4=self.use_mxfp4, + w1_scale=w1_scale, + w2_scale=w2_scale, + block_size=None, + a1_scale=None, + a2_scale=None, + w1_bias=w1_bias, + w2_bias=w2_bias, + alpha=use_alpha, + limit=use_limit, + is_vnni=getattr(self, "packed_weight", False), + ) + + # Restore original shape + output = output.view(in_shape) + + return output, expert_weights + + +__all__ = ["fused_moe_cpp", "MegaBlocksMoeMLP"] diff --git a/build/torch214-cxx11-cu126-x86_64-linux/grouped_gemm/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/grouped_gemm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b91c8308f0c24f4c4171b6e4f15b6f76dabf295a --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/grouped_gemm/__init__.py @@ -0,0 +1,2 @@ +from . import ops +from . import backend diff --git a/build/torch214-cxx11-cu126-x86_64-linux/grouped_gemm/backend.py b/build/torch214-cxx11-cu126-x86_64-linux/grouped_gemm/backend.py new file mode 100644 index 0000000000000000000000000000000000000000..c7ef28ced79c830dae934177f059c1f4ddc24aad --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/grouped_gemm/backend.py @@ -0,0 +1,33 @@ +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# On ROCm there is no CUTLASS grouped GEMM; dispatch to the vendored AITER +# Triton kernels instead. On CUDA we use the compiled CUTLASS `gmm` op. +_IS_ROCM = torch.version.hip is not None + +if _IS_ROCM: + from .._grouped_gemm_triton import adapter as backend +else: + # We import the backend operations from the megablocks package as + # grouped_gemm is vendored in megablocks in this repository. + from .._ops import ops as backend # type: ignore + +def _allocate_output(a, b, batch_sizes, trans_a, trans_b): + assert not (trans_a and trans_b) + assert batch_sizes.ndim == 1, "Expected 1d tensor for batch_sizes" + assert a.ndim == 2, "Expected 2d tensor for 'a'" + assert b.ndim == (2 if trans_a else 3) + + shape = ( + (batch_sizes.shape[0], a.shape[1], b.shape[1]) + if trans_a else + (a.shape[0], (b.shape[1] if trans_b else b.shape[2])) + ) + return torch.empty(*shape, device=a.device, dtype=a.dtype) + +def gmm(a, b, batch_sizes, trans_a=False, trans_b=False, c=None): + if c is None: + c = _allocate_output(a, b, batch_sizes, trans_a, trans_b) + backend.gmm(a, b, c, batch_sizes, trans_a, trans_b) + return c diff --git a/build/torch214-cxx11-cu126-x86_64-linux/grouped_gemm/ops.py b/build/torch214-cxx11-cu126-x86_64-linux/grouped_gemm/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..4b30dd14e23837ea3b12334f4e31337ed9ad2b69 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/grouped_gemm/ops.py @@ -0,0 +1,33 @@ +from . import backend +import torch + + +class GroupedGemm(torch.autograd.Function): + + @staticmethod + def forward(ctx, a, b, batch_sizes, trans_b): + ctx.save_for_backward(a, b, batch_sizes) + ctx.trans_b = trans_b + return backend.gmm(a, b, batch_sizes, trans_a=False, trans_b=trans_b) + + @staticmethod + def backward(ctx, grad): + grad = grad.contiguous() + a, b, batch_sizes = ctx.saved_tensors + trans_b = ctx.trans_b + + agrad = None + if ctx.needs_input_grad[0]: + agrad = backend.gmm( + grad, b, batch_sizes, trans_a=False, trans_b=not trans_b) + + bgrad = None + if ctx.needs_input_grad[1]: + lhs, rhs = (grad, a) if trans_b else (a, grad) + bgrad = backend.gmm( + lhs, rhs, batch_sizes, trans_a=True, trans_b=False) + return agrad, bgrad, None, None + + +def gmm(a, b, batch_sizes, trans_b=False): + return GroupedGemm.apply(a, b, batch_sizes, trans_b) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/grouped_gemm_util.py b/build/torch214-cxx11-cu126-x86_64-linux/grouped_gemm_util.py new file mode 100644 index 0000000000000000000000000000000000000000..1d6d49fc46b0a57ad46e4179df3cc1ac2a24f7ae --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/grouped_gemm_util.py @@ -0,0 +1,31 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +import warnings + +_grouped_gemm_is_available: bool = False +try: + # import grouped_gemm + pass + _grouped_gemm_is_available = True +except ImportError as error: + warnings.warn('Grouped GEMM not available.') + + +def grouped_gemm_is_available(): + return _grouped_gemm_is_available + + +def assert_grouped_gemm_is_available(): + msg = ( + 'Grouped GEMM not available. Please run ' + '`pip install git+https://github.com/tgale96/grouped_gemm@main`.', + ) + assert _grouped_gemm_is_available, msg + + +# backend = grouped_gemm.backend if grouped_gemm_is_available() else None +# ops = grouped_gemm.ops if grouped_gemm_is_available() else None + + +from .grouped_gemm import backend as ops +from .grouped_gemm import ops as backend diff --git a/build/torch214-cxx11-cu126-x86_64-linux/layers.py b/build/torch214-cxx11-cu126-x86_64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..7c91edcd9e2d1a4ef9eac90217ff481f08ab1886 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/layers.py @@ -0,0 +1,1232 @@ +import torch +import torch.distributed as dist + +from typing import Optional, Any, TYPE_CHECKING + +from . import _layers +from . import ops + +# Conditional import for meta kernel registration +if TYPE_CHECKING: + + def register_fake(fn): + return lambda name: fn + +else: + try: + from torch.library import register_fake + except ImportError: + try: + from torch.library import impl_abstract as register_fake + except ImportError: + # Fallback for older PyTorch versions + def register_fake(op_name): + def decorator(fn): + return fn + + return decorator + + +# Meta kernel implementations for torch.compile compatibility +def _install_meta_kernels(): + """Install meta kernels for existing MegaBlocks operations""" + + # Create wrapper functions that check for compilation and return meta tensors + + # Patch ops.sort + if hasattr(ops, "sort"): + original_sort = ops.sort + + def sort_with_meta(x, end_bit=None): + if torch.compiler.is_compiling(): + # print("Using meta kernel for sort") + # Meta implementation - return tensors with correct shape/dtype/device + return torch.empty_like(x), torch.empty_like(x) + # print("Using original sort kernel") + return original_sort(x, end_bit) + + ops.sort = sort_with_meta + + # Patch ops.histogram + if hasattr(ops, "histogram"): + original_histogram = ops.histogram + + def histogram_with_meta(x, max_val): + if torch.compiler.is_compiling(): + # Meta implementation + return torch.empty((max_val,), dtype=torch.int32, device=x.device) + return original_histogram(x, max_val) + + ops.histogram = histogram_with_meta + + # Patch ops.inclusive_cumsum + if hasattr(ops, "inclusive_cumsum"): + original_inclusive_cumsum = ops.inclusive_cumsum + + def inclusive_cumsum_with_meta(x, dim): + if torch.compiler.is_compiling(): + # Meta implementation + return torch.empty_like(x) + return original_inclusive_cumsum(x, dim) + + ops.inclusive_cumsum = inclusive_cumsum_with_meta + + # Patch ops.binned_gather + if hasattr(ops, "binned_gather"): + original_binned_gather = ops.binned_gather + + def binned_gather_with_meta(x, indices, bins, bin_size, top_k): + if torch.compiler.is_compiling(): + # Meta implementation - output shape based on bin_size + if x.dim() >= 2: + hidden_size = x.size(-1) + return torch.empty( + (bin_size, x.size(1), hidden_size), + dtype=x.dtype, + device=x.device, + ) + else: + return torch.empty((bin_size,), dtype=x.dtype, device=x.device) + return original_binned_gather(x, indices, bins, bin_size, top_k) + + ops.binned_gather = binned_gather_with_meta + + # Patch ops.binned_scatter + if hasattr(ops, "binned_scatter"): + original_binned_scatter = ops.binned_scatter + + def binned_scatter_with_meta(x, indices, weights, bins, top_k): + if torch.compiler.is_compiling(): + # Meta implementation - typically reduces to 2D + if x.dim() >= 3: + return torch.empty( + (x.size(1), x.size(2)), dtype=x.dtype, device=x.device + ) + else: + return torch.empty_like(x) + return original_binned_scatter(x, indices, weights, bins, top_k) + + ops.binned_scatter = binned_scatter_with_meta + + # Patch ops.gather + if hasattr(ops, "gather"): + original_gather = ops.gather + + def gather_with_meta(x, indices, bin_ids, bins, top_k): + if torch.compiler.is_compiling(): + # Meta implementation + if x.dim() >= 2: + hidden_size = x.size(-1) + return torch.empty( + (indices.numel(), hidden_size), dtype=x.dtype, device=x.device + ) + else: + return torch.empty(indices.shape, dtype=x.dtype, device=x.device) + return original_gather(x, indices, bin_ids, bins, top_k) + + ops.gather = gather_with_meta + + # Patch ops.scatter + if hasattr(ops, "scatter"): + original_scatter = ops.scatter + + def scatter_with_meta(x, indices, bin_ids, weights, bins, top_k): + if torch.compiler.is_compiling(): + # Meta implementation - restore sequence shape + seq_len = ( + indices.size(0) // top_k + if indices.numel() > 0 and top_k > 0 + else x.size(0) + ) + if x.dim() >= 2: + return torch.empty( + (seq_len, x.size(-1)), dtype=x.dtype, device=x.device + ) + else: + return torch.empty((seq_len,), dtype=x.dtype, device=x.device) + return original_scatter(x, indices, bin_ids, weights, bins, top_k) + + ops.scatter = scatter_with_meta + + # Patch ops.replicate + if hasattr(ops, "replicate"): + original_replicate = ops.replicate + + def replicate_with_meta(x, bins, num_outputs): + if torch.compiler.is_compiling(): + # Meta implementation + return torch.empty( + (x.shape[0], num_outputs), dtype=x.dtype, device=x.device + ) + return original_replicate(x, bins, num_outputs) + + ops.replicate = replicate_with_meta + + # Patch ops.repeat (if it's a regular function) + if hasattr(ops, "repeat"): + original_repeat = ops.repeat + + def repeat_with_meta(x, repeats): + if torch.compiler.is_compiling(): + # Meta implementation + if isinstance(repeats, (tuple, list)): + new_shape = list(x.shape) + for i, rep in enumerate(repeats): + if i < len(new_shape): + new_shape[i] *= rep + return torch.empty(new_shape, dtype=x.dtype, device=x.device) + else: + new_shape = [x.size(0) * repeats] + list(x.shape[1:]) + return torch.empty(new_shape, dtype=x.dtype, device=x.device) + return original_repeat(x, repeats) + + ops.repeat = repeat_with_meta + + +# Install meta kernels on import +try: + _install_meta_kernels() +except Exception as e: + # If meta kernel installation fails, continue without them + # torch.compile may not work but the library will still function + import warnings + + warnings.warn( + f"Failed to install meta kernels for torch.compile support: {e}", UserWarning + ) + + +# Set the expert model parallel attributes on a tensor +def set_expert_model_parallel_attributes( + tensor: torch.Tensor, + is_parallel: bool, +): + assert not hasattr(tensor, "expert_model_parallel") + setattr(tensor, "expert_model_parallel", is_parallel) + + +# Get the expert model parallel attributes from a tensor +def expert_sharding_degree( + world_size: int, + moe_num_experts: int, +) -> int: + esd = min(world_size, moe_num_experts) + if (moe_num_experts % esd) != 0: + raise ValueError(f"Cannot shard {moe_num_experts} experts {esd} ways.") + return esd + + +# Calculate the hidden sharding degree based on world size and expert sharding degree +def hidden_sharding_degree( + world_size: int, + moe_num_experts: int, + ffn_hidden_size: int, +) -> int: + esd = expert_sharding_degree(world_size, moe_num_experts) + hsd = world_size // esd + if (ffn_hidden_size % hsd) != 0: + raise ValueError(f"Cannot shard {ffn_hidden_size} features {hsd} ways.") + if (esd * hsd) != world_size: + raise ValueError( + f"Invalid sharding. expert_sharding_degree ({esd}) * hidden_sharding_degree ({hsd}) != world_size ({world_size})." + ) + return hsd + + +# Calculate the number of experts per rank based on world size and expert sharding degree +def experts_per_rank( + moe_num_experts: int, + world_size: int, +) -> int: + return moe_num_experts // expert_sharding_degree(world_size, moe_num_experts) + + +# Calculate the number of features per rank based on ffn hidden size and hidden sharding degree +def features_per_rank( + ffn_hidden_size: int, world_size: int, moe_num_experts: int +) -> int: + return ffn_hidden_size // hidden_sharding_degree( + world_size, moe_num_experts, ffn_hidden_size + ) + + +# Apply jitter to the input tensor +def apply_jitter(x: torch.Tensor, moe_jitter_eps: float) -> torch.Tensor: + low = 1.0 - moe_jitter_eps + high = 1.0 + moe_jitter_eps + noise = torch.rand(x.size(), dtype=x.dtype, device=x.device) + return x * (low + noise * (high - low)) + + +# Compute the top-k scores from the logits +def compute_top_k(scores: torch.Tensor, moe_top_k: int): + if moe_top_k == 1: + return scores.max(dim=-1, keepdim=True) + return torch.topk(scores, moe_top_k, dim=-1) + + +# Route tokens to experts and compute expert weights and indices +def route_tokens( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + moe_top_k: int, + moe_num_experts: int, + moe_jitter_eps: float = None, + moe_normalize_expert_weights: int = None, + uniform_expert_assignment: bool = False, + training: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if training and moe_jitter_eps is not None: + x = apply_jitter(x, moe_jitter_eps) + + x_flat = x.view(-1, x.shape[-1]) + logits = torch.nn.functional.linear(x_flat, router_weight, router_bias) + expert_weights, expert_indices = compute_top_k(logits, moe_top_k) + expert_weights = expert_weights.softmax(dim=-1) + if moe_normalize_expert_weights is not None: + expert_weights = expert_weights / torch.norm( + expert_weights, + p=moe_normalize_expert_weights, + dim=-1, + keepdim=True, + ) + if uniform_expert_assignment: + expert_indices = _layers.router._uniform_expert_assignment( + expert_indices, + moe_num_experts, + ) + + return logits, expert_weights, expert_indices + + +# Scale the gradient of the weights +def scale_grad( + w: torch.Tensor, + gradient_scale: Optional[float] = None, +) -> torch.Tensor: + if gradient_scale is None: + return w + return _layers.mlp.scale_gradient(w, gradient_scale) + + +# Forward pass for the MLP layer +def mlp_forward( + x: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w1_bias: torch.Tensor, + w2_bias: torch.Tensor, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + limit: float = 7.0, +): + # Scale weights + w1 = scale_grad(w1, gradient_scale) + w2 = scale_grad(w2, gradient_scale) + w1_bias = scale_grad(w1_bias, gradient_scale) + w2_bias = scale_grad(w2_bias, gradient_scale) + + # Resolve dtensors + w1 = _layers.mlp.resolve_dtensor(w1) + w2 = _layers.mlp.resolve_dtensor(w2) + w1_bias = _layers.mlp.resolve_dtensor(w1_bias) + w2_bias = _layers.mlp.resolve_dtensor(w2_bias) + + # Forward pass + gate_up = torch.bmm(x, w1) + w1_bias[..., None, :] + gate, up = gate_up[..., ::2], gate_up[..., 1::2] + gate = gate.clamp(min=None, max=limit) + up = up.clamp(min=-limit, max=limit) + glu = gate * torch.sigmoid(gate * alpha) + next_states = torch.bmm(((up + 1) * glu), w2) + next_states += w2_bias[..., None, :] + return next_states + +# Shared expert MLP forward pass +def shared_mlp_forward( + x: torch.Tensor, + up_proj_weight: torch.Tensor, + down_proj_weight: torch.Tensor, + up_proj_bias: Optional[torch.Tensor] = None, + down_proj_bias: Optional[torch.Tensor] = None, + activation_fn: Optional[Any] = None, + gradient_scale: Optional[float] = None, +) -> torch.Tensor: + # Default activation function + if activation_fn is None: + activation_fn = torch.nn.functional.gelu + + # Scale weights + up_proj_weight = scale_grad(up_proj_weight, gradient_scale) + down_proj_weight = scale_grad(down_proj_weight, gradient_scale) + if up_proj_bias is not None: + up_proj_bias = scale_grad(up_proj_bias, gradient_scale) + if down_proj_bias is not None: + down_proj_bias = scale_grad(down_proj_bias, gradient_scale) + + # Resolve dtensors + up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight) + down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight) + if up_proj_bias is not None: + up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias) + if down_proj_bias is not None: + down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias) + + # Up projection + x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias) + + # Activation + x = activation_fn(x) + + # Down projection + x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias) + + return x + + +# Combine outputs from shared expert and regular experts +def combine_expert_shared_outputs( + shared_expert_out: torch.Tensor, + expert_out: torch.Tensor, + shared_expert_weighted_sum: bool = False, + moe_top_k: int = 1, +) -> torch.Tensor: + if shared_expert_weighted_sum: + # Weighted sum based on number of experts used + total_experts = moe_top_k + 1 + shared_weight = 1.0 / total_experts + expert_weight = moe_top_k / total_experts + return shared_expert_out * shared_weight + expert_out * expert_weight + else: + # Simple addition + return shared_expert_out + expert_out + + +# Global variable to store load balancing loss +_LOAD_BALANCING_LOSS = [] + + +def save_load_balancing_loss(loss): + global _LOAD_BALANCING_LOSS + _LOAD_BALANCING_LOSS.append(loss) + + +def get_load_balancing_loss(): + global _LOAD_BALANCING_LOSS + return _LOAD_BALANCING_LOSS + + +def clear_load_balancing_loss(): + global _LOAD_BALANCING_LOSS + _LOAD_BALANCING_LOSS.clear() + + +def batched_load_balancing_loss(args): + if args.moe_loss_weight == 0: + return 0.0 + + tokens_per_expert, expert_scores = zip(*get_load_balancing_loss()) + num_layers_per_pipeline_stage = args.num_layers // args.pipeline_model_parallel_size + if args.num_layers_per_virtual_pipeline_stage is not None: + num_layers_per_pipeline_stage = args.num_layers_per_virtual_pipeline_stage + + if len(tokens_per_expert) != num_layers_per_pipeline_stage: + raise ValueError( + f"Expected {num_layers_per_pipeline_stage} token_per_experts " + f"but found {len(tokens_per_expert)}.\nnum_layers = " + f"{args.num_layers}\npipeline_model_parallel_size = " + f"{args.pipeline_model_parallel_size}\n" + "num_layers_per_virtual_pipeline_stage" + f" = {args.num_layers_per_virtual_pipeline_stage}", + ) + if len(expert_scores) != num_layers_per_pipeline_stage: + raise ValueError( + f"Expected {num_layers_per_pipeline_stage} expert_scores " + f"but found {len(tokens_per_expert)}.\nnum_layers = " + f"{args.num_layers}\npipeline_model_parallel_size = " + f"{args.pipeline_model_parallel_size}\n" + "num_layers_per_virtual_pipeline_stage" + f" = {args.num_layers_per_virtual_pipeline_stage}", + ) + + # Verify the shape of the tokens_per_expert and expert_scores tensors. + assert all( + (x.ndim == 1 and x.numel() == args.moe_num_experts for x in tokens_per_expert) + ) + + tokens = expert_scores[0].shape[0] + assert all( + ( + ( + x.ndim == 2 + and x.shape[1] == args.moe_num_experts + and x.shape[0] == tokens + ) + for x in expert_scores + ) + ) + + # Concatenate the contributions of each layer and convert to + # the correct types and formats for the dot product. + expert_scores = torch.cat(expert_scores, dim=1) + if args.moe_lbl_in_fp32: + expert_scores = expert_scores.float() + if tokens != 0: + expert_scores = expert_scores.mean(dim=0) + else: + expert_scores = expert_scores.sum(dim=0) + tokens_per_expert = torch.cat(tokens_per_expert).to(expert_scores.dtype) + + expected_values = num_layers_per_pipeline_stage * args.moe_num_experts + assert tokens_per_expert.numel() == expected_values + assert expert_scores.numel() == expected_values + + # Calculate the total scale across all factors. + # + # loss_weight * num_experts / (num_layers * tokens * top_k) + scale_numerator = args.moe_num_experts * args.moe_loss_weight + scale_denominator = args.num_layers * tokens * args.moe_top_k + scale = scale_numerator / scale_denominator + return scale * torch.dot(tokens_per_expert, expert_scores) + + +# Calculate the expert capacity based on tokens, top_k, number of experts, +# expert parallel group, capacity factor, and whether expert model parallelism is used. +def expert_capacity( + tokens: int, + top_k: int, + num_experts: int, + expert_parallel_group: int, + moe_capacity_factor: float, + moe_expert_model_parallelism: bool, +) -> int: + world_size = ( + dist.get_world_size(expert_parallel_group) + if moe_expert_model_parallelism + else 1 + ) + + tokens_per_expert = top_k * tokens * world_size / num_experts + return int(moe_capacity_factor * tokens_per_expert) + + +def load_balancing_loss( + tokens_per_expert: torch.Tensor, + expert_scores: torch.Tensor, + top_k: int, + num_experts: int, +): + assert len(expert_scores.size()) == 2 + tokens, num_experts = expert_scores.size() + assert num_experts == num_experts + assert len(tokens_per_expert.size()) == 1 + (num_experts,) = tokens_per_expert.size() + assert num_experts == num_experts + scale = num_experts / (tokens * top_k) + return scale * torch.dot( + tokens_per_expert.to(expert_scores.dtype), + expert_scores.mean(dim=0), + ) + + +def indices_and_bins( + top_expert: torch.Tensor, + sort_end_bit: int, + num_experts: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + top_expert = top_expert.int() + + # Ensure contiguous memory layout + top_expert = top_expert.contiguous() + + # Ensure CUB knows which device to use + with torch.cuda.device(top_expert.device): + output = ops.sort(top_expert, sort_end_bit) + bin_ids, indices = output + tokens_per_expert = ops.histogram(top_expert, num_experts) + bins = ops.inclusive_cumsum(tokens_per_expert, 0) + + bins = bins.view(1) if not len(bins.size()) else bins + return indices, bin_ids, bins, tokens_per_expert + + +def expert_capacity_fn( + tokens: int, + top_k: int, + num_experts: int, + expert_parallel_group: torch.distributed.ProcessGroup, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = False, +) -> int: + world_size = ( + dist.get_world_size(expert_parallel_group) + if moe_expert_model_parallelism + else 1 + ) + tokens_per_expert = top_k * tokens * world_size / num_experts + return int(moe_capacity_factor * tokens_per_expert) + + +def permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capacity, + top_k, + w1, + w2, + w1_bias, + w2_bias, + gradient_scale, + alpha, +): + # Route tokens to experts + x = x.view(-1, x.shape[-1]) + + # Ensure CUB knows which device to use + with torch.cuda.device(x.device): + x = ops.binned_gather(x, indices, bins, expert_capacity, top_k) + + # Expert computation + x = mlp_forward(x, w1, w2, w1_bias, w2_bias, gradient_scale, alpha) + + # Ensure CUB knows which device to use + with torch.cuda.device(x.device): + # Route tokens back + out = ops.binned_scatter(x, indices, expert_weights, bins, top_k) + return out + + +def forward_once( + x: torch.Tensor, + expert_weights: torch.Tensor, + top_experts: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w1_bias: torch.Tensor, + w2_bias: torch.Tensor, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + sort_end_bit: int = 0, + top_k: int = 4, + num_experts: int = 128, + expert_parallel_group: int = None, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = False, + mlp_impl: Optional[str] = None, +): + # x: [sl, bs, hs] + # expert_weights: [sl * bs, top-k] + # top_experts: [sl * bs, top-k] + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + + with torch.no_grad(): + indices, bin_ids, bins, tokens_per_expert = indices_and_bins( + top_experts, sort_end_bit, num_experts + ) + + # Calculate expert capacity + sl, bs, _ = x.size() + + expert_capacity = expert_capacity_fn( + sl * bs, + top_k, + num_experts, + expert_parallel_group, + moe_capacity_factor, + moe_expert_model_parallelism, + ) + + if expert_capacity == 0: + expert_capacity = torch.max(tokens_per_expert).item() + + x = permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capacity, + top_k, + w1, + w2, + w1_bias, + w2_bias, + gradient_scale, + alpha, + ) + return x, tokens_per_expert + + +def parallel_forward_once( + x: torch.Tensor, + expert_weights: torch.Tensor, + top_experts: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w1_bias: torch.Tensor, + w2_bias: torch.Tensor, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + sort_end_bit: int = 0, + top_k: int = 4, + num_experts: int = 128, + expert_parallel_group: torch.distributed.ProcessGroup = None, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = True, + hidden_size: int = 1152, + mlp_impl: Optional[str] = "grouped", +): + # Flatten inputs + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + + # TODO: remove debugging var + # my_rank = dist.get_rank(expert_parallel_group) if expert_parallel_group else 0 + + with torch.no_grad(): + # Step 1: Local permutation setup + indices, bin_ids, bins, tokens_per_expert = indices_and_bins( + top_experts, sort_end_bit, num_experts + ) + + # Calculate sharding parameters + world_size = dist.get_world_size(expert_parallel_group) + hidden_sharding_deg = hidden_sharding_degree( + world_size, num_experts, hidden_size + ) + experts_per_rank_val = experts_per_rank(num_experts, world_size) + + # Replicate token counts for hidden sharding + repeated_tokens_per_expert = ops.repeat( + tokens_per_expert, (hidden_sharding_deg,) + ) + + # Exchange token counts across devices + parallel_tokens_per_expert = torch.empty_like(repeated_tokens_per_expert) + + # Ensure CUB knows which device to use + tpe_handle = dist.all_to_all_single( + parallel_tokens_per_expert, + repeated_tokens_per_expert, + group=expert_parallel_group, + async_op=True, + ) + + # Step 2: Local permutation - group tokens by target device + x = x.view(-1, x.shape[-1]) # [sl * bs, hs] + x = ops.gather(x, indices, bin_ids, bins, top_k) + + # Step 3: Compute communication counts and exchange tokens + with torch.no_grad(): + tpe_handle.wait() + + # Reshape for per-device calculations + repeated_tokens_per_expert = repeated_tokens_per_expert.view( + world_size, experts_per_rank_val + ) + parallel_tokens_per_expert = parallel_tokens_per_expert.view( + world_size, experts_per_rank_val + ) + + # Calculate send/recv counts + send_counts = repeated_tokens_per_expert.cpu().sum(dim=-1).tolist() + # recv_counts = parallel_tokens_per_expert.cpu().sum(dim=-1).tolist() + parallel_tokens_per_expert_cpu = parallel_tokens_per_expert.cpu() + recv_counts = parallel_tokens_per_expert_cpu.sum(dim=-1).tolist() + tokens_received = sum(recv_counts) + + # Replicate for hidden sharding + x = ops.repeat(x, (hidden_sharding_deg, 1)) + + # Cross-device token exchange + parallel_x, parallel_x_handle = _layers.all_to_all.all_to_all( + x, recv_counts, send_counts, expert_parallel_group, async_op=True + ) + + with torch.no_grad(): + # Step 4: Setup for local expert computation + replicate_bins = ops.inclusive_cumsum(parallel_tokens_per_expert.flatten(), 0) + replicate_bins = ( + replicate_bins.view(1) if not len(replicate_bins.size()) else replicate_bins + ) + + # Create expert indices for received tokens + parallel_top_expert = torch.remainder( + torch.arange( + num_experts * hidden_sharding_deg, + dtype=torch.int32, + device=indices.device, + ), + experts_per_rank_val, + ) + parallel_top_expert = ops.replicate( + parallel_top_expert.unsqueeze(dim=0), + replicate_bins, + tokens_received, + ).flatten() + + # Sort tokens by expert assignment + parallel_bin_ids, parallel_indices = ops.sort( + parallel_top_expert, + sort_end_bit, + ) + + # Calculate bins for local experts + parallel_tokens_per_expert = parallel_tokens_per_expert.sum( + dim=0, dtype=torch.int + ) + parallel_bins = ops.inclusive_cumsum(parallel_tokens_per_expert, 0) + parallel_bins = ( + parallel_bins.view(1) if not len(parallel_bins.size()) else parallel_bins + ) + + # Calculate expert capacity + expert_capacity = expert_capacity_fn( + tokens_received, + top_k, + experts_per_rank_val, + expert_parallel_group, + moe_capacity_factor, + moe_expert_model_parallelism, + ) + if expert_capacity == 0: + expert_capacity = torch.max(parallel_tokens_per_expert).item() + + # Locally permute the tokens and perform the expert computation. + # Block to make sure that the cross-device permutation is complete. + if mlp_impl == "grouped": + # GroupedMLP requires counts on CPU. We can use the tensor already + # moved to CPU for the prior all_to_all, which avoids an extra + # device synchronization. + parallel_tokens_per_expert = parallel_tokens_per_expert_cpu.sum( + dim=0, + dtype=torch.int, + ) + + # Step 5: Expert computation + parallel_x_handle.wait() + + parallel_x = permute_and_compute( + parallel_x, + parallel_tokens_per_expert, + parallel_indices, + parallel_bin_ids, + None, # expert_weights + parallel_bins, + expert_capacity, + top_k=1, + w1=w1, + w2=w2, + w1_bias=w1_bias, + w2_bias=w2_bias, + gradient_scale=gradient_scale, + alpha=alpha, + ) + + # Step 6: Reverse communication - send results back + x, _ = _layers.all_to_all.all_to_all( + parallel_x, send_counts, recv_counts, expert_parallel_group + ) + + # Step 7: Reduce across hidden sharding dimension + shape = (hidden_sharding_deg, -1, hidden_size) + x = x.view(shape).sum(dim=0) + + # Step 8: Final local unpermutation + x = ops.scatter(x, indices, bin_ids, expert_weights, bins, top_k) + + return x, tokens_per_expert.flatten() + + +def moe_forward( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: Optional[torch.Tensor], + moe_top_k: int, + moe_num_experts: int, + moe_jitter_eps: float = None, + moe_normalize_expert_weights: int = None, + uniform_expert_assignment: bool = False, + training: bool = False, + w1: torch.Tensor = None, + w2: torch.Tensor = None, + w1_bias: torch.Tensor = None, + w2_bias: torch.Tensor = None, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + sort_end_bit: int = 0, + expert_parallel_group: torch.distributed.ProcessGroup = None, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = False, + forward_fn: Any = None, + hidden_size: int = None, + mlp_impl: str = "grouped", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + # Route tokens to experts + logits, expert_weights, expert_indices = route_tokens( + x, + router_weight, + router_bias, + moe_top_k, + moe_num_experts, + moe_jitter_eps, + moe_normalize_expert_weights, + uniform_expert_assignment, + training, + ) + + # Create router scores for output + router_scores = ( + torch.zeros_like(logits) + .scatter_(1, expert_indices, expert_weights) + .transpose(0, 1) + ) + + in_shape = x.size() + + # Prepare forward function arguments + forward_args = { + "x": x, + "expert_weights": expert_weights, + "top_experts": expert_indices, + "w1": w1, + "w2": w2, + "w1_bias": w1_bias, + "w2_bias": w2_bias, + "gradient_scale": gradient_scale, + "alpha": alpha, + "sort_end_bit": sort_end_bit, + "top_k": moe_top_k, + "num_experts": moe_num_experts, + "expert_parallel_group": expert_parallel_group, + "moe_capacity_factor": moe_capacity_factor, + "moe_expert_model_parallelism": moe_expert_model_parallelism, + "mlp_impl": mlp_impl, + } + + # Add hidden_size for parallel forward + if moe_expert_model_parallelism and hidden_size is not None: + forward_args["hidden_size"] = hidden_size + elif moe_expert_model_parallelism and hidden_size is None: + # Infer hidden_size from input shape + forward_args["hidden_size"] = x.shape[-1] + + # Compute expert outputs + x, tokens_per_expert = forward_fn(**forward_args) + + # Save load balancing loss if needed + moe_loss_weight = 0.0 # Can be made configurable + if training and moe_loss_weight > 0: + save_load_balancing_loss((tokens_per_expert, logits)) + + # Restore original shape + x = x.view(in_shape) + + return x, expert_weights, router_scores + + +def moe_forward_with_shared_expert( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: Optional[torch.Tensor], + moe_top_k: int, + moe_num_experts: int, + moe_jitter_eps: float = None, + moe_normalize_expert_weights: int = None, + uniform_expert_assignment: bool = False, + training: bool = False, + w1: torch.Tensor = None, + w2: torch.Tensor = None, + w1_bias: torch.Tensor = None, + w2_bias: torch.Tensor = None, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + sort_end_bit: int = 0, + expert_parallel_group: torch.distributed.ProcessGroup = None, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = False, + forward_fn: Any = None, + hidden_size: int = None, + mlp_impl: str = "grouped", + # Shared expert parameters + shared_up_proj_weight: Optional[torch.Tensor] = None, + shared_down_proj_weight: Optional[torch.Tensor] = None, + shared_up_proj_bias: Optional[torch.Tensor] = None, + shared_down_proj_bias: Optional[torch.Tensor] = None, + shared_expert_weighted_sum: bool = False, + shared_activation_fn: Optional[Any] = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + # First, compute regular MoE forward pass + expert_out, expert_weights, router_scores = moe_forward( + x=x, + router_weight=router_weight, + router_bias=router_bias, + moe_top_k=moe_top_k, + moe_num_experts=moe_num_experts, + moe_jitter_eps=moe_jitter_eps, + moe_normalize_expert_weights=moe_normalize_expert_weights, + uniform_expert_assignment=uniform_expert_assignment, + training=training, + w1=w1, + w2=w2, + w1_bias=w1_bias, + w2_bias=w2_bias, + gradient_scale=gradient_scale, + alpha=alpha, + sort_end_bit=sort_end_bit, + expert_parallel_group=expert_parallel_group, + moe_capacity_factor=moe_capacity_factor, + moe_expert_model_parallelism=moe_expert_model_parallelism, + forward_fn=forward_fn, + hidden_size=hidden_size, + mlp_impl=mlp_impl, + ) + + # If shared expert weights provided, compute shared expert output + if shared_up_proj_weight is not None and shared_down_proj_weight is not None: + shared_expert_out = shared_mlp_forward( + x=x, + up_proj_weight=shared_up_proj_weight, + down_proj_weight=shared_down_proj_weight, + up_proj_bias=shared_up_proj_bias, + down_proj_bias=shared_down_proj_bias, + activation_fn=shared_activation_fn, + gradient_scale=gradient_scale, + ) + + # Combine expert outputs + combined_out = combine_expert_shared_outputs( + shared_expert_out=shared_expert_out, + expert_out=expert_out, + shared_expert_weighted_sum=shared_expert_weighted_sum, + moe_top_k=moe_top_k, + ) + + return combined_out, expert_weights, router_scores + + # Return regular MoE output if no shared expert + return expert_out, expert_weights, router_scores + + +def create_shared_expert_weights( + hidden_size: int, + shared_expert_hidden_size: int, + device: torch.device, + dtype: torch.dtype, + init_method: Any, + output_layer_init_method: Any = None, +) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + + if output_layer_init_method is None: + output_layer_init_method = init_method + + # Create weight tensors + up_proj_weight = torch.empty( + shared_expert_hidden_size, + hidden_size, + device=device, + dtype=dtype, + ) + down_proj_weight = torch.empty( + hidden_size, + shared_expert_hidden_size, + device=device, + dtype=dtype, + ) + + # Initialize weights + init_method(up_proj_weight) + output_layer_init_method(down_proj_weight) + + # No bias by default + return up_proj_weight, down_proj_weight, None, None + + +# HACK: Extract device_mesh from pre-hook closure - required for transformers integration +# This exists because device_mesh is trapped in hook closures with no model attribute +# Fragile - breaks if hook structure changes or Python internals change +# TODO: Replace with a more robust solution when available +def get_device_mesh(model): + # Extract device_mesh from child's unused pre_hook closure + try: + # Find the pre-hook that contains 'device_mesh' in its closure + hook = next( + h + for h in model.experts._forward_pre_hooks.values() + if "device_mesh" in h.__code__.co_freevars + ) + # Extract the device_mesh from the closure + return hook.__closure__[ + hook.__code__.co_freevars.index("device_mesh") + ].cell_contents + except Exception: + return None + + +class MegaBlocksMoeMLP(torch.nn.Module): + can_torch_compile: bool = True + + def forward(self, x: torch.Tensor) -> torch.Tensor: + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + gradient_scale = getattr(self.experts, "gradient_scale", None) + alpha = getattr(self.experts, "alpha", 1.0) + moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0) + moe_jitter_eps = getattr(self.experts, "jitter_eps", None) + moe_normalize_expert_weights = getattr( + self.experts, "normalize_expert_weights", None + ) + uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False) + + expert_parallel_group = getattr(self, "expert_parallel_group", None) + if expert_parallel_group is None: + device_mesh = get_device_mesh(self) + expert_parallel_group = device_mesh.get_group() if device_mesh else None + + has_parallel = ( + expert_parallel_group is not None + and dist.is_initialized() + and dist.get_world_size(expert_parallel_group) > 1 + ) + forward_fn = parallel_forward_once if has_parallel else forward_once + + sort_end_bit = max( + int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1 + ) + mlp_impl = getattr(self, "mlp_impl", "grouped") + output, expert_weights_out, *_ = moe_forward( + x=x, + router_weight=self.router.weight, + router_bias=self.router.bias, + moe_top_k=moe_top_k, + moe_num_experts=moe_num_experts, + moe_jitter_eps=moe_jitter_eps, + moe_normalize_expert_weights=moe_normalize_expert_weights, + uniform_expert_assignment=uniform_expert_assignment, + training=self.training, + w1=self.experts.gate_up_proj, + w2=self.experts.down_proj, + w1_bias=self.experts.gate_up_proj_bias, + w2_bias=self.experts.down_proj_bias, + gradient_scale=gradient_scale, + alpha=alpha, + sort_end_bit=sort_end_bit, + expert_parallel_group=expert_parallel_group, + moe_capacity_factor=moe_capacity_factor, + moe_expert_model_parallelism=has_parallel, + forward_fn=forward_fn, + hidden_size=self.experts.hidden_size, + mlp_impl=mlp_impl, + ) + return output, expert_weights_out + + +# Export main classes +__all__ = ["MegaBlocksMoeMLP", "MegaBlocksMoeMLPWithSharedExpert"] + + +class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP): + + def __init__(self): + super().__init__() + # Shared expert weights will be set by the user + self.shared_up_proj_weight = None + self.shared_down_proj_weight = None + self.shared_up_proj_bias = None + self.shared_down_proj_bias = None + self.shared_expert_weighted_sum = False + self.shared_activation_fn = None + + def set_shared_expert_weights( + self, + up_proj_weight: torch.Tensor, + down_proj_weight: torch.Tensor, + up_proj_bias: Optional[torch.Tensor] = None, + down_proj_bias: Optional[torch.Tensor] = None, + weighted_sum: bool = False, + activation_fn: Optional[Any] = None, + ): + self.shared_up_proj_weight = up_proj_weight + self.shared_down_proj_weight = down_proj_weight + self.shared_up_proj_bias = up_proj_bias + self.shared_down_proj_bias = down_proj_bias + self.shared_expert_weighted_sum = weighted_sum + self.shared_activation_fn = activation_fn + + def forward(self, x: torch.Tensor) -> torch.Tensor: + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + gradient_scale = getattr(self.experts, "gradient_scale", None) + alpha = getattr(self.experts, "alpha", 1.0) + moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0) + moe_jitter_eps = getattr(self.experts, "jitter_eps", None) + moe_normalize_expert_weights = getattr( + self.experts, "normalize_expert_weights", None + ) + uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False) + + expert_parallel_group = getattr(self, "expert_parallel_group", None) + if expert_parallel_group is None: + device_mesh = get_device_mesh(self) + expert_parallel_group = device_mesh.get_group() if device_mesh else None + + has_parallel = ( + expert_parallel_group is not None + and dist.is_initialized() + and dist.get_world_size(expert_parallel_group) > 1 + ) + forward_fn = parallel_forward_once if has_parallel else forward_once + + sort_end_bit = max( + int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1 + ) + mlp_impl = getattr(self, "mlp_impl", "grouped") + + output, expert_weights_out, *_ = moe_forward_with_shared_expert( + x=x, + router_weight=self.router.weight, + router_bias=self.router.bias, + moe_top_k=moe_top_k, + moe_num_experts=moe_num_experts, + moe_jitter_eps=moe_jitter_eps, + moe_normalize_expert_weights=moe_normalize_expert_weights, + uniform_expert_assignment=uniform_expert_assignment, + training=self.training, + w1=self.experts.gate_up_proj, + w2=self.experts.down_proj, + w1_bias=self.experts.gate_up_proj_bias, + w2_bias=self.experts.down_proj_bias, + gradient_scale=gradient_scale, + alpha=alpha, + sort_end_bit=sort_end_bit, + expert_parallel_group=expert_parallel_group, + moe_capacity_factor=moe_capacity_factor, + moe_expert_model_parallelism=has_parallel, + forward_fn=forward_fn, + hidden_size=self.experts.hidden_size, + mlp_impl=mlp_impl, + # Shared expert parameters + shared_up_proj_weight=self.shared_up_proj_weight, + shared_down_proj_weight=self.shared_down_proj_weight, + shared_up_proj_bias=self.shared_up_proj_bias, + shared_down_proj_bias=self.shared_down_proj_bias, + shared_expert_weighted_sum=self.shared_expert_weighted_sum, + shared_activation_fn=self.shared_activation_fn, + ) + return output, expert_weights_out + + +# Patch for XPU or CPU support +if hasattr(torch, "xpu") and torch.xpu.is_available(): + from .xpu_fused_moe import MegaBlocksMoeMLP + +from .cpu_moe_cpp import CPUMegaBlocksMoeMLP diff --git a/build/torch214-cxx11-cu126-x86_64-linux/metadata.json b/build/torch214-cxx11-cu126-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..0050ff179520e825b1508f5d25e9de42e3431f4c --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/metadata.json @@ -0,0 +1,117 @@ +{ + "name": "megablocks", + "id": "_megablocks_cuda_addf474", + "version": 1, + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "7.0", + "7.2", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0+PTX" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "o7Y9XLXezcJogNTNoC4HMHsd2ZZ3CTwzYftkUthOtno=", + "_grouped_gemm_triton/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "_grouped_gemm_triton/_triton_kernels/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "_grouped_gemm_triton/_triton_kernels/gmm.py": "11SvNqMt/LDCgjJBOmZlHg9yn9NbGw1UYKzNR/JKtcQ=", + "_grouped_gemm_triton/adapter.py": "tP00AypBJn4rDqrBDMMwBb6MwJ6s9dg18GNQxpLiVoA=", + "_grouped_gemm_triton/configs.py": "aHXlKELHIy8J82m1kClQRxDkad4LCL8BOeEpq4/IqKQ=", + "_grouped_gemm_triton/gmm.py": "1iqK7N3vsahgQ4bbHsOMA8hoR5uQhh1PmbFWjOLSKDo=", + "_grouped_gemm_triton/utils/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "_grouped_gemm_triton/utils/_triton/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "_grouped_gemm_triton/utils/_triton/arch_info.py": "Vo7tziqn9GLA8S4xrKT2555aKFxo75Oja4ek2zNReuY=", + "_grouped_gemm_triton/utils/_triton/pid_preprocessing.py": "Lx1S7li2xArbeql/pEFfyul7yU61agwuYHTsZ7jG0os=", + "_grouped_gemm_triton/utils/gmm_common.py": "pQMeVpP9kk/bixwJE309iceJPa8H7mZV8b6+OtbQvl8=", + "_grouped_gemm_triton/utils/logger.py": "CXuJZCBzQDTrW80OyJ2IuhbHhayLy12Ake1KGxiuYIA=", + "_layers/__init__.py": "bwEDLk8kbLuT1BWZPPDVUiDsWsJ3wC3nI+O+dSGAoSU=", + "_layers/activation_fn.py": "lJrZcw+1oLw0XOyjYb6AfmxX1n6g2Snbph6H3p5D0oU=", + "_layers/all_to_all.py": "SZMHtoTy6B76caEE0shoK1y4iHqWp+lWqdR/30d9oBQ=", + "_layers/arguments.py": "wYFjDwBidH5MjMr4cZGDEcRvW1IgdLQvcYt/qt5S0dk=", + "_layers/common.py": "8m+bNabk+jB+C8o797GdNpGzx8m3nbnG66FQdnDZRVI=", + "_layers/dmlp_registry.py": "z7iqmxiJx9eACdaHOCmnY9EGKV/zePhqcf8FVboaoK8=", + "_layers/dmoe.py": "bOaQ0PFOOLu45HxJNRFuKPB8VedUzsduekkqJV6WrQ8=", + "_layers/gelu.py": "fcq4OBCSCcPjCk7PYZydIpbdn/iUyMtL2TqhonHqt7k=", + "_layers/glu.py": "D27/lTdF05N6TdSPoyW1MWgTGyaXJlUCo5uos54zabU=", + "_layers/memory_test.py": "4jm72yP6WBV5zelWQGJQIHHZLPisFEtar+moOSnxXMk=", + "_layers/mlp.py": "xx51A1pGeBt/B8m7pI5hjhVqhEllpxvldulXYzFmA3Y=", + "_layers/moe.py": "vVqEx4Ppgw+QbLHOJo5uSI/qXDikA99ZGApPOcBJtmg=", + "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=", + "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=", + "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=", + "_megablocks_cuda_addf474.abi3.so": "VvLqh7ruAZpHEKNBJOlG0V+II5KOhtA2a0gHgOfoVCo=", + "_ops.py": "WXlV4XwtKQjdQ6nap3k9TalaHkK0ypeSTshOrV6Hczo=", + "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=", + "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=", + "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=", + "benchmark_util.py": "WFsGWGaJvWvwdx4w33WCfNzSJkVuQrZmZCaAodJsws8=", + "cpu_fused_moe.py": "XN2eQBFaihbJrZ/cu9jS8h5dcoiUNh82gjJQUARcBYA=", + "cpu_moe_cpp.py": "dvv4Ynn/FWObs421L4hAdYZPiPs2BS/sbEOuMMIUPVM=", + "grouped_gemm/__init__.py": "fxW9amtzy/Z0V/OXSYMmibD7XvpYBGEv81KMxvocFKY=", + "grouped_gemm/backend.py": "Yd//ihpvu/pxEoGBq3iD1zMggnYv6quwZMdg1R4kUEI=", + "grouped_gemm/ops.py": "Kj7yRgv7afB5SVwXIkN9RNIF2kxD2qrZ0dVsKX5ph2s=", + "grouped_gemm_util.py": "bVPrdsEtBVJJ+v/bNrGNf/QHNlVL4VxHg92HvaL3oq4=", + "layers.py": "mNTlX2oKJyHBJllRd24lPNdGLa2H9cEUgVLvkZJpoNo=", + "ops/__init__.py": "MHQsiMEgkKeZCa8TAYFMGDFnwFIgijfAN7DhMUt/48U=", + "ops/all_to_all_benchmark.py": "sxXSCSDdsh+62IWpKNTyBphLTngY1jlgJ1WgNooSEO8=", + "ops/binned_gather.py": "pkWne+ThfiTR/bMbDs2QDimKxAI4klqD1FdvhytE7lU=", + "ops/binned_scatter.py": "5NxocRmgYDkTAntDtXBlF4c2r8rfbLVnBRPA6f4/L4U=", + "ops/cumsum.py": "upuWnnp125kQhJrIMLtG09cI1zzxEy06P5wYhOwbpI4=", + "ops/gather.py": "a8j/btP1VCdJ7F7FC58MFPxFuH9tUhjMuQnEzwjAgGw=", + "ops/histogram.py": "v5+pHF+/fjKQAkNLIsHyljMUakMTdW5GzA094PgRoGk=", + "ops/histogram_benchmark.py": "oX8VolN+/PVCWaPHkY0k7SNn4QpadHi80Sla/3NScNI=", + "ops/matmul_benchmark.py": "SZ3L/CatnT4J5PyOU/Ii/b5y5fw3HNEaR/xY6PjrJ8A=", + "ops/padded_gather.py": "VtPH3Ge1jodeug5MWwXnxnmGIi0D1MOOaXi2CcU8lpI=", + "ops/padded_scatter.py": "UO1BwsaezK3AikD8CHoJOEjmC4h+H5fcOX/bvO8FWXw=", + "ops/padded_scatter_benchmark.py": "w34qlIc3Nn0AgB2k7iIFqQXgDNplMVqCr5i8GrgrvZ0=", + "ops/permute_benchmark.py": "ISYPkVZknTnfB8fhUVloSyfEOe7ApKQKtNPvEq3P9qg=", + "ops/repeat.py": "j9l1A/4LVhcFuSAaneFlImq7RGNXie0zfu2nMlEhHI4=", + "ops/replicate.py": "NN5843jJz0OdfYS6YaPCpRD8dGk66S3SHTzR51PrwUI=", + "ops/round_up.py": "/CQSCYmjW0Mrv8TFdfZs/ZpXHHz18b3E4aLHvmfUxOY=", + "ops/scatter.py": "YgM76I/dYAzRNN+jzAdaRWan56GiwuCZRR7xy3RAY8Y=", + "ops/sort.py": "35D3oG4AT0Jd6+awwfoL+GmkcNtdRLno7LcSFidqKts=", + "ops/sort_benchmark.py": "Bo985Xsxq/F0PMU92xKCo+srSKVdheQLey2GqFlwZXY=", + "ops/stk_autocast.py": "SQArC9BFDXVManccbEWibSg0E+/R3NQlkp0cB4mIjXo=", + "ops/sum.py": "fcjWM5Aiz1wQ+UYBN2Y/OFILsiGKSLgmlkeqfVRccD4=", + "ops/topology.py": "+lfjfxDjB+st3K8wZnRY5fEanSIWT+3SJncx9bu+6P0=", + "stk/__init__.py": "QFugfJQ5q1Wdn/+VQyfs1YKns79i3FiywiLevg0ymXo=", + "stk/backend/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "stk/backend/autocast.py": "ppk7mL/WQUHaLqEgsW3h5ZGhldC8WyIKmaIS6kAe4lA=", + "stk/backend/sputnik.py": "Mpbo0M1ES39PdaIo65ef3/mQ8XHtLbO3ayqufaIZLiw=", + "stk/backend/triton_kernels.py": "RCthwX2feppkqBK3pvm+/37nhF89lLHgP12c6xrfxXU=", + "stk/matrix.py": "Xxo1RwOtKFgVXExhlAoIC67V+eWKZRClDdbBfzUgfy0=", + "stk/ops/__init__.py": "Vqbt6C0qWEPaWypsCoaHGoZRbdPcFrLg1ahlAlp4l+Y=", + "stk/ops/eltwise_ops.py": "O9TqOMvZ2NZ7KNFaeuCmJOfg1Hz5ad8rvYwbGAH7/+0=", + "stk/ops/eltwise_ops_test.py": "rD1oPPu+vXHlyJnEJOpM9L5wdCeiDbn98CbCH8mQ0wE=", + "stk/ops/linear_ops.py": "ZWOpLTXMdw/ZTyxL4W74CSSZQaUKvtfiz3ev0OnaVQ0=", + "stk/ops/linear_ops_test.py": "kl8nMX1EAvZ6lF7oN4U4bW2z0hera0YxqTPiNljLmzk=", + "stk/ops/matrix_ops.py": "+hhUk7IZZFVVzIk1pttWFFSOtKwlGEBVR2V4eyL27K4=", + "stk/ops/matrix_ops_test.py": "Bb4fk4+IzNC0zOFq0Tz/JwNaLWyhtWEDgEUzyVdbnrE=", + "stk/random/__init__.py": "Gr/+Gl20jJ6fJVRfoAUxEcCPeChILQzrzhKVVAWQsRI=", + "stk/random/random_ops.py": "AluIDj3Vo2ReKAsZz07yydbtndDbJoglekRGIaCoRZE=", + "stk/random/random_ops_test.py": "B+a4Zb+OPhBtpVA48oc09sBL/AAQO9j/YUFI0nnXxzk=", + "xpu_fused_moe.py": "sPmuRSpkH1Bmk4hdpz53us8Qpc6yinT9rkld3T8WreI=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "commit": "7d2828e6592de641b44fb8eb896719101a2ab2fb", + "dirty": false + }, + "kernel": { + "commit": "addf4741fafb55b4c101b15985dfaeea7c7c2f3a", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch214-cxx11-cu126-x86_64-linux/metadata.json.sigstore b/build/torch214-cxx11-cu126-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..449048e404732800c45705613abf19b2df0b978a --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUKeMemPCJyuShgh3MhMXdIZ+DS3swCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI3MDcxNDE0WhcNMjYwODI3MDcyNDE0WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAErHaunB+83zdTzIrn/y+vWf+9M8rOvN4hg+kUsaRMYhYdyGWowHCMdev8boaGh5Fq3iLXORHIVuBxQNr96ZTXTqOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUDtLMZ4+RQV1YF/1EB+y4DptjoAEwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGFkZGY0NzQxZmFmYjU1YjRjMTAxYjE1OTg1ZGZhZWVhN2M3YzJmM2EwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzMwNDUzOTU5NTIvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoEIRwNEAAAQDAEcwRQIhAOJNaw4CejDUU8TPgIk1WmmRZOhBKGOZIVywN1yUMSQHAiBvcpo5EcAKOhFhHfVC3TkuPOiIEPmZt01a4nehpfhnIDAKBggqhkjOPQQDAwNoADBlAjAr3oZJILSAyH1HwcbEFKSXJsGTbm6jG9LZslXtH9Usre5Y9FYUn3mU1+boBqVACwcCMQDMIBYsA9doNcIz8ZC0wPZhfgvvjycpTrjDP0IUELTxnqqiHfflmFT5DtNv7PRsP9w="}, "tlogEntries":[{"logIndex":"2614049034", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787814855", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQCOwH91gnzrt1ub/lsEBUATaqlI83X2u9mz6FGrVRXXhwIhAKacwHDQzlGlGfXPjyAmE94S0Kg74ANJ6Y5HYFcAZJT3"}, "inclusionProof":{"logIndex":"2492144772", "rootHash":"5uAYUoZYEIF456mWg7woGu/iDyJzUrlhRjxy4/z0D9s=", "treeSize":"2492144796", "hashes":["N6iMsq0fRf6RNs+yUDZXlCfESatQLRyoD5zQ26FawdY=", "LadlGQ2qe8mUTwcBc3SWTt6MWi7v39/1X2iqmOGDa5M=", "w6u/gBOybsj8YsRgzgg+RMV7aDXAz2OhvZ3JAGESIsc=", "pVop5RkpuwUwEPA7Vaa8K8Sdp02F+gE+pztkqSIuSJU=", "yd3xfzkUitDRAY229+rq/bQHrvX0pfdC6b/zaZjgJaY=", "tX6xfx7l6K10NEfEZHsTZneVXcVhfdWkODMuNY+ukac=", "JbYm9/m/JN8Zi+pgvvB/Pwv6te3HLQLa0r0nBi4DSeE=", "VFlTxnDU4W+k/WA3ARTvzCiqVt1yxQIMUyUzMCugJ2Q=", "CFNMzmJ8N+uNx37P2asMqvrW9tw4axAVHnXLRO39tFs=", "NypqAbB+mNix4l1ahGeP9fRWiQzZCm4P1664makjjq4=", "8SoaeA7IeJPoIiveQbBfHCWXSDQabIk2PW+FT49NSEQ=", "VGrFt5JvlIUQvoGEdBBxrCe3lYbu581OhgLq/Xdyat4=", "7TfDCaMWgEwSJEZd/e33b24919rF4IDJLmbqNNRNhZs=", "0HfdxDt/zGugZLIuHasrdLEW9s8OaTykSSqBEDFQXdw=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2492144796\n5uAYUoZYEIF456mWg7woGu/iDyJzUrlhRjxy4/z0D9s=\n\n— rekor.sigstore.dev wNI9ajBFAiEAgVrVTx2UOqSMvWvCp69xAPQ3q9UK/8L3AwxewWAvvlQCICnMIF9am66QPHAU8rd7RtM3MzstFiw/ILelQHBQtVAN\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI5M2Q1Y2MyNmU4ZmEwNzQwOWM4MDExZTEyMTdjZjFmMjQ5NjBjYTg3NDZhNWYyNjJhMDU4ZTMxNTJmYzYyNTA3In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUURjbTVnYUJaVXc5V1RzeCt4SlhwTER3d3lFaDNOcmlVbitaY1VhVDRpN3NRSWdNYlNtdEY5L2tqOEpyM0F1OUxpcEc2emV0VnZyQjY2UllNaElTV3U3YjFFPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlMyVk5aVzFRUTBwNWRWTm9aMmd6VFdoTldHUkpXaXRFVXpOemQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1ROTlJHTjRUa1JGTUZkb1kwNU5hbGwzVDBSSk0wMUVZM2xPUkVVd1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ5U0dGMWJrSXJPRE42WkZSNlNYSnVMM2tyZGxkbUt6bE5PSEpQZGs0MGFHY3JhMVVLYzJGU1RWbG9XV1I1UjFkdmQwaERUV1JsZGpoaWIyRkhhRFZHY1ROcFRGaFBVa2hKVm5WQ2VGRk9jamsyV2xSWVZIRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZFZEV4TkNsbzBLMUpSVmpGWlJpOHhSVUlyZVRSRWNIUnFiMEZGZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVpWMUpyV21wUk0wNUVSbTFaVjFwcFRsUldhVTVIVFhoTlJFWnBDazFVVlRWUFJGWnJXbTFHYkZwWFJUTlplbVJxVFcxWmVsbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWxOVkZVMVQwUldhMXB0Um13S1dsZEZNMWw2WkdwTmJWbDZXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWtLVFZSVk5VOUVWbXRhYlVac1dsZEZNMWw2WkdwTmJWbDZXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZEdhMXBIV1RBS1RucFJlRnB0Um0xWmFsVXhXV3BTYWsxVVFYaFpha1V4VDFSbk1WcEhXbWhhVjFab1RqSk5NMWw2U20xTk1rVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWsxM1RrUlZlazlVVlRWT1ZFbDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUZTVkozVGtWQlFVRlJSQXBCUldOM1VsRkphRUZQU2s1aGR6UkRaV3BFVlZVNFZGQm5TV3N4VjIxdFVscFBhRUpMUjA5YVNWWjVkMDR4ZVZWTlUxRklRV2xDZG1Od2J6VkZZMEZMQ2s5b1JtaElabFpETTFScmRWQlBhVWxGVUcxYWREQXhZVFJ1Wldod1ptaHVTVVJCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEJjak52V2tvS1NVeFRRWGxJTVVoM1kySkZSa3RUV0VwelIxUmliVFpxUnpsTVduTnNXSFJJT1ZWemNtVTFXVGxHV1ZWdU0yMVZNU3RpYjBKeFZrRkRkMk5EVFZGRVRRcEpRbGx6UVRsa2IwNWpTWG80V2tNd2QxQmFhR1puZG5acWVXTndWSEpxUkZBd1NWVkZURlI0Ym5GeGFVaG1abXh0UmxRMVJIUk9kamRRVW5OUU9YYzlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgx1Cn7W17NYEGk9c7uwA8oLhw3OYGyMwXabeGS4mTa98CFQCYP/UqoPCj7ESoK7pudzFhKXxntBgPMjAyNjA4MjcwNzE0MTRaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgyNzA3MTQxNFowLwYJKoZIhvcNAQkEMSIEIM5sStFvrZlsSucnwPkVVEcQ0wn1975Rc0YY8YcIlhxJMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQDWUD0IQoCa57BlR6xM7B4erIAcEoyPfxixX89cjOIAauI0mpor3A+6a1YGS1mnG50CMDO+3Pm8dcblphyC+RTKTwR5lN4RWM2og7VdFJYC4FHHan+7hv9v6Trfxdtm7dEAIw=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"k9XMJuj6B0CcgBHhIXzx8klgyodGpfJioFjjFS/GJQc="}, "signature":"MEUCIQDcm5gaBZUw9WTsx+xJXpLDwwyEh3NriUn+ZcUaT4i7sQIgMbSmtF9/kj8Jr3Au9LipG6zetVvrB66RYMhISWu7b1E="}} \ No newline at end of file diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b944080df810d0b0cfc571f3009b0098a651f9b7 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/__init__.py @@ -0,0 +1,35 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from .binned_gather import binned_gather +from .binned_scatter import binned_scatter +from .cumsum import exclusive_cumsum, inclusive_cumsum +from .gather import gather +from .histogram import histogram +from .padded_gather import padded_gather +from .padded_scatter import padded_scatter +from .repeat import repeat +from .replicate import replicate +from .round_up import round_up +from .scatter import scatter +from .sort import sort +from .sum import sum +from .topology import topology + +__all__ = [ + 'binned_gather', + 'binned_scatter', + 'exclusive_cumsum', + 'inclusive_cumsum', + 'gather', + 'histogram', + 'padded_gather', + 'padded_scatter', + 'repeat', + 'replicate', + 'round_up', + 'scatter', + 'sort', + 'sum', + 'topology', +] diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/all_to_all_benchmark.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/all_to_all_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..4c939818edca3345f6344bbc7cef07ffe3cd0181 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/all_to_all_benchmark.py @@ -0,0 +1,63 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.distributed as dist + +# from megablocks import benchmark_util +# from megablocks.layers.all_to_all import all_to_all + +from .. import benchmark_util +from .._layers.all_to_all import all_to_all + +_ALL_TO_ALL_BENCHMARK = ( + (8, 1024), + (16, 1024), + (32, 1024), + (64, 1024), + (128, 1024), + (256, 1024), + (512, 1024), + (1024, 1024), + (2 * 1024, 1024), + (4 * 1024, 1024), + (8 * 1024, 1024), + (16 * 1024, 1024), + (32 * 1024, 1024), + (64 * 1024, 1024), + (128 * 1024, 1024), + (256 * 1024, 1024), + (512 * 1024, 1024), + (1024 * 1024, 1024), +) + + +def benchmark_all_to_all(group, sl, hs): + world_size = dist.get_world_size(group) + assert (sl % world_size) == 0 + send_recv_sizes = [sl // world_size] * world_size + + x = torch.randn((sl, hs)).cuda().half() + + details = { + 'world_size': world_size, + 'message_size (B)': send_recv_sizes[0] * hs * 2, # 2B elements. + } + + def benchmark(): + return all_to_all(x, send_recv_sizes, send_recv_sizes, group) + + time, std = benchmark_util.benchmark_function(benchmark) + + if dist.get_rank(group) == 0: + benchmark_util.log_benchmark('All-To-All', details, time, std) + + +if __name__ == '__main__': + assert dist.is_available() + group = dist.init_process_group(backend='nccl') + local_rank = dist.get_rank(group) + torch.cuda.set_device(local_rank) + + for args in _ALL_TO_ALL_BENCHMARK: + benchmark_all_to_all(group, *args) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/binned_gather.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/binned_gather.py new file mode 100644 index 0000000000000000000000000000000000000000..189a7fa3518d660f29ea32e7a04827164af98d60 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/binned_gather.py @@ -0,0 +1,37 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for binned_gather kernel. +class BinnedGatherOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bins: torch.Tensor, + bin_size: int, + top_k: int, + ): + ctx.save_for_backward(indices, bins) + ctx.top_k = top_k + return kernels.binned_gather(x, indices, None, bins, bin_size, top_k) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + indices, bins = ctx.saved_tensors + out = kernels.binned_scatter(grad, indices, None, bins, ctx.top_k) + return out, None, None, None, None + + +binned_gather = BinnedGatherOp.apply diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/binned_scatter.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/binned_scatter.py new file mode 100644 index 0000000000000000000000000000000000000000..cb937c0c106662ce8108c1cb926f8f063b163d3d --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/binned_scatter.py @@ -0,0 +1,59 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for binned_scatter kernel. +class BinnedScatterOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + top_k: int, + ): + assert len(x.size()) == 3 + ctx.bin_size = x.size(1) + ctx.top_k = top_k + + # TODO(tgale): Don't save 'x' for backwards if we don't need to + # calculate the gradient w.r.t. 'weights'. + ctx.save_for_backward(x, indices, weights, bins) + return kernels.binned_scatter(x, indices, weights, bins, top_k) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + x, indices, weights, bins = ctx.saved_tensors + out = kernels.binned_gather( + grad, + indices, + weights, + bins, + ctx.bin_size, + ctx.top_k, + ) + + wgrad = None + if ctx.needs_input_grad[2]: + wgrad = kernels.binned_scatter_wgrad( + x, + grad, + indices, + bins, + ctx.top_k, + ) + return out, None, wgrad, None, None + + +binned_scatter = BinnedScatterOp.apply diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/cumsum.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/cumsum.py new file mode 100644 index 0000000000000000000000000000000000000000..e2b7572391e20045d335cf7337246e8a9b9f57ef --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/cumsum.py @@ -0,0 +1,52 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + # import megablocks_ops as ops # type: ignore + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + + +# Autograd wrappers for cumsum kernels. +# NOTE: Does not support gradients. +class ExclusiveCumsumOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, dim: int): + if len(x.size()) == 1: + x = x.view([1, -1]) + out = torch.empty_like(x) + ops.exclusive_cumsum(x, 1, out) + return out.squeeze() + out = torch.empty_like(x) + ops.exclusive_cumsum(x, dim, out) + return out + + +exclusive_cumsum = ExclusiveCumsumOp.apply + + +class InclusiveCumsumOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, dim: int) -> torch.Tensor: + if len(x.size()) == 1: + x = x.view([1, -1]) + out = torch.empty_like(x) + ops.inclusive_cumsum(x, 1, out) + return out.squeeze() + out = torch.empty_like(x) + ops.inclusive_cumsum(x, dim, out) + return out + + +inclusive_cumsum = InclusiveCumsumOp.apply diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/gather.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/gather.py new file mode 100644 index 0000000000000000000000000000000000000000..f1f87c1e7bed8d3589dd790805234976e0b05898 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/gather.py @@ -0,0 +1,38 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for gather kernel. +class GatherOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + bins: torch.Tensor, + top_k: int, + ): + ctx.save_for_backward(indices, bin_ids, bins) + ctx.top_k = top_k + return kernels.gather(x, indices, bin_ids, None, bins, top_k) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + + indices, bin_ids, bins = ctx.saved_tensors + out = kernels.scatter(grad, indices, bin_ids, None, bins, ctx.top_k) + return out, None, None, None, None, None + + +gather = GatherOp.apply diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/histogram.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/histogram.py new file mode 100644 index 0000000000000000000000000000000000000000..7b3f058ec373cbba7555704fb5e4212c3cc75d9d --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/histogram.py @@ -0,0 +1,27 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + + +# Autograd wrapper for histogram kernel. +# NOTE: Does not support gradients. +class HistogramOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, max_val: float): + return ops.histogram(x, max_val) + + +histogram = HistogramOp.apply diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/histogram_benchmark.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/histogram_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..a0b9c6047567b87a295979498142230d1b0c9411 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/histogram_benchmark.py @@ -0,0 +1,78 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import numpy as np +import torch +# from absl.testing import parameterized + +from .. import ops + +_HISTOGRAM_TESTS = ( + (16384, torch.int32, 2), + (16384, torch.int32, 4), + (16384, torch.int32, 8), + (16384, torch.int32, 16), + (16384, torch.int32, 32), + (16384, torch.int32, 64), + (16384, torch.int32, 128), + (16384, torch.int32, 256), +) + + +def benchmark_function(fn, iterations=10): + # Run once to get rid of startup overhead. + fn() + times = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + times = np.array(times) + return times.mean(), times.std(), times.max(), times.min() + + +def log_benchmark(arguments, mean_t, std_t): + print('=' * 60) + print('Benchmark Parameters:') + for (key, value) in arguments.items(): + print(f'{key} = {value}') + print('Results:') + print('mean / std = {:.2f}ms / {:.2f}ms'.format(mean_t, std_t)) + print('=' * 60) + + +# class HistogramBenchmark(parameterized.TestCase): +# +# @parameterized.parameters(*_HISTOGRAM_TESTS) +# def testHistogram(self, n, dtype, max_val): +# x = torch.randint(0, max_val, (n,)).cuda().to(dtype) +# +# mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.histogram(x, max_val),) +# arguments = { +# 'n': n, +# 'dtype': dtype, +# 'max_val': max_val, +# } +# log_benchmark(arguments, mean_t, std_t) +# +# @parameterized.parameters(*_HISTOGRAM_TESTS) +# def testTorchHistogram(self, n, dtype, max_val): +# x = torch.randint(0, 128, (n,)).cuda().to(dtype) +# +# mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.histc(x, max_val, 0, max_val - 1),) +# arguments = { +# 'n': n, +# 'dtype': dtype, +# 'max_val': max_val, +# } +# log_benchmark(arguments, mean_t, std_t) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/matmul_benchmark.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/matmul_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..4c4b9b8866ffed2eb769b77f2320c82e5491ae0e --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/matmul_benchmark.py @@ -0,0 +1,415 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + + +# import stk + +# try: +# import stk +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/ops/matmul_benchmark.py is needed.', +# ) + +from .. import stk + +import torch +# from absl.testing import parameterized + +from .. import benchmark_util, ops + + +# Calling tensor.t() calls tensor.transpose(0, 1) which calls +# torch.as_strided(...). Circumvent this chain to avoid an overhead +# this adds. +def transpose_view(x): + return torch.as_strided( + x, + (x.shape[1], x.shape[0]), + (x.stride()[1], x.stride()[0]), + ) + + +_MATMUL_TESTS = ( + (64 * 1024, 512, 2048, 64), + (32 * 1024, 768, 3072, 64), + (8 * 1024, 1024, 4096, 64), + (4 * 2048, 4096, 4 * 4096, 4), +) + + +def log_benchmark(name, arguments, time, std, flops): + benchmark_util.log_benchmark(name, arguments, time, std) + print('flops = {:.2f}B'.format(flops / 1e9)) + print('throughput = {:.2f}T'.format(flops / 1e9 / time)) + print('=' * 60) + + +# class MatmulBenchmark(parameterized.TestCase): +# +# def build_sparse_matrix(self, x, padded_bins, fhs, ne): +# blocking = 128 +# padded_tokens, _ = x.size() +# assert padded_tokens % blocking == 0 +# assert fhs % blocking == 0 +# +# # Offsets for the sparse matrix. All rows have the +# # same number of nonzero blocks dictated by the +# # dimensionality of a single expert. +# block_rows = padded_tokens // blocking +# blocks_per_row = fhs // blocking +# offsets = torch.arange( +# 0, +# block_rows * blocks_per_row + 1, +# blocks_per_row, +# dtype=torch.int32, +# device=x.device, +# ) +# +# # Indices for the sparse matrix. The indices for +# # the intermediate matrix are dynamic depending +# # on the mapping of tokens to experts. +# column_indices = ops.topology( +# padded_bins, +# blocking, +# block_rows, +# blocks_per_row, +# ) +# data = torch.empty( +# column_indices.numel(), +# blocking, +# blocking, +# dtype=torch.float16, +# device=x.device, +# ) +# shape = (padded_tokens, fhs * ne) +# row_indices = stk.ops.row_indices(shape, data, offsets, column_indices) +# return stk.Matrix(shape, data, row_indices, column_indices, offsets) +# +# def build_input_matrix(self, sl, hs, ne): +# x = torch.randn((sl, hs)).cuda().half() +# +# # Assign tokens to experts uniformly. +# top_expert = torch.arange(0, sl).cuda().int() % ne +# +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(top_expert, ne) +# padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) +# padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# out = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, 1) +# return out, padded_bins +# +# def build_weight_matrix(self, ne, hs, fhs): +# return torch.randn((hs, ne * fhs)).cuda().half() +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_Fwd_SDD_NT(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# topo = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# w = transpose_view(w) +# +# def benchmark(): +# return stk.ops.sdd(x, w, topo) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0::Fwd::SDD::NT', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_GradX_DSD_NN(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# topo = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# +# def benchmark(): +# return stk.ops.dsd(topo, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0::GradX::DSD::NN', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_GradW_DSD_TN(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# topo = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# topo = topo.t() +# +# def benchmark(): +# return stk.ops.dsd(topo, x) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0::GradW::DSD::TN', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_Fwd_DSD_NN(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# x = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# +# def benchmark(): +# return stk.ops.dsd(x, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::Fwd::DSD::NN', +# arguments, +# mean_t, +# std_t, +# x.nnz * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_GradX_SDD_NT(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# x = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# out = stk.ops.dsd(x, w) +# w = transpose_view(w) +# +# def benchmark(): +# return stk.ops.sdd(out, w, x) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::GradX::SDD::NT', +# arguments, +# mean_t, +# std_t, +# x.nnz * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_GradW_DSD_TN(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# x = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# out = stk.ops.dsd(x, w) +# x = x.t() +# +# def benchmark(): +# return stk.ops.dsd(x, out) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::GradW::DSD::TN', +# arguments, +# mean_t, +# std_t, +# x.nnz * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_Fwd_DDD_NT(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, hs)).cuda().half() +# w = torch.randn((ne, hs, fhs)).cuda().half() +# +# w = w.transpose(1, 2).contiguous() +# w = w.transpose(1, 2) +# +# def benchmark(): +# return torch.bmm(x, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0::Fwd:DDD::NT', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_GradX_DDD_NN(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, hs)).cuda().half() +# w = torch.randn((ne, hs, fhs)).cuda().half() +# out = torch.bmm(x, w) +# w = w.transpose(1, 2).contiguous() +# +# def benchmark(): +# return torch.bmm(out, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0:GradX:DDD::NN', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_GradW_DDD_TN(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, hs)).cuda().half() +# w = torch.randn((ne, hs, fhs)).cuda().half() +# out = torch.bmm(x, w) +# out = out.transpose(1, 2) +# +# def benchmark(): +# return torch.bmm(out, x) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0:GradW:DDD::TN', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_Fwd_DDD_NN(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, fhs)).cuda().half() +# w = torch.randn((ne, fhs, hs)).cuda().half() +# +# def benchmark(): +# return torch.bmm(x, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::Fwd::DDD::NN', +# arguments, +# mean_t, +# std_t, +# x.numel() * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_GradX_DDD_NT(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, fhs)).cuda().half() +# w = torch.randn((ne, fhs, hs)).cuda().half() +# out = torch.bmm(x, w) +# w = torch.transpose(w, 1, 2) +# +# def benchmark(): +# return torch.bmm(out, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::GradX::DDD::NT', +# arguments, +# mean_t, +# std_t, +# x.numel() * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_GradW_DDD_TN(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, fhs)).cuda().half() +# w = torch.randn((ne, fhs, hs)).cuda().half() +# out = torch.bmm(x, w) +# x = torch.transpose(x, 1, 2) +# +# def benchmark(): +# return torch.bmm(x, out) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::GradW::DDD::TN', +# arguments, +# mean_t, +# std_t, +# x.numel() * hs * 2, +# ) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/padded_gather.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/padded_gather.py new file mode 100644 index 0000000000000000000000000000000000000000..c1cf4047c9494394d2a3884ba8830179013db7ff --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/padded_gather.py @@ -0,0 +1,55 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for padded_gather kernel. +class PaddedGatherOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + bins: torch.Tensor, + padded_bins: torch.Tensor, + top_k: int, + ): + ctx.save_for_backward(indices, bin_ids, bins, padded_bins) + ctx.top_k = top_k + return kernels.padded_gather( + x, + indices, + bin_ids, + None, + bins, + padded_bins, + top_k, + ) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + + indices, bin_ids, bins, padded_bins = ctx.saved_tensors + out = kernels.padded_scatter( + grad, + indices, + bin_ids, + None, + bins, + padded_bins, + ctx.top_k, + ) + return out, None, None, None, None, None + + +padded_gather = PaddedGatherOp.apply diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/padded_scatter.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/padded_scatter.py new file mode 100644 index 0000000000000000000000000000000000000000..61e021b81497e472cda5d72bdac557a0ca92d262 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/padded_scatter.py @@ -0,0 +1,98 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for padded_scatter kernel. +class PaddedScatterOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + padded_bins: torch.Tensor, + top_k: int, + ): + maybe_x = [x] if ctx.needs_input_grad[3] else [] + ctx.save_for_backward( + indices, + bin_ids, + weights, + bins, + padded_bins, + *maybe_x, + ) + ctx.top_k = top_k + ctx.x_shape = x.shape + return kernels.padded_scatter( + x, + indices, + bin_ids, + weights, + bins, + padded_bins, + top_k, + ) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + saved_tensors = ctx.saved_tensors + + indices, bin_ids, weights, bins, padded_bins = saved_tensors[:5] + dgrad = None + if ctx.needs_input_grad[0]: + dgrad = kernels.padded_gather( + grad, + indices, + bin_ids, + weights, + bins, + padded_bins, + ctx.top_k, + ) + + wgrad = None + if ctx.needs_input_grad[3]: # need wgrad + x = saved_tensors[-1] + wgrad = kernels.padded_scatter_wgrad( + x, + grad, + indices, + bin_ids, + bins, + padded_bins, + ctx.top_k, + ) + return dgrad, None, None, wgrad, None, None, None, None + + +def padded_scatter( + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + padded_bins: torch.Tensor, + top_k: int, +): + return PaddedScatterOp.apply( + x, + indices, + bin_ids, + weights, + bins, + padded_bins, + top_k, + ) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/padded_scatter_benchmark.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/padded_scatter_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..3cbe4735891446b46f93170c64c23fe63632bf93 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/padded_scatter_benchmark.py @@ -0,0 +1,66 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import torch +# from absl.testing import parameterized + +from .. import benchmark_util, ops + +_PADDED_SCATTER_BENCHMARK = ( + # dMoE-Medium, 8-way EMP. + (1024 * 16, 1024, 8, 4), + # dMoE-Medium, post-all-to-all. + (1024 * 16 * 4, 1024, 8, 1), +) + + +# class PaddedScatterTest(parameterized.TestCase): +# +# @parameterized.parameters(*_PADDED_SCATTER_BENCHMARK) +# def testPaddedScatter(self, sl, hs, ne, top_k): +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# +# # Randomly assign tokens to experts. +# top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(top_expert, ne) +# padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) +# padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# +# # Sample weights for the scatter reduce. +# weights = torch.rand((sl * top_k,)).cuda().half() +# +# # Gather the data to prepare for backwards. +# x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k) +# +# def benchmark(): +# return ops.padded_scatter( +# x, +# indices, +# bin_ids, +# weights, +# bins, +# padded_bins, +# top_k, +# ) +# +# time, std = benchmark_util.benchmark_function(benchmark) +# benchmark_util.log_benchmark( +# 'Padded Scatter', +# { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# 'top_k': top_k, +# }, +# time, +# std, +# ) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/permute_benchmark.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/permute_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..697abddbb3a2082ec4ddd6d94f89f7faabb34b40 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/permute_benchmark.py @@ -0,0 +1,149 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import torch +# from absl.testing import parameterized + +from .. import benchmark_util, ops + +_PERMUTE_TESTS = ( + (16384, 768, 2), + (16384, 768, 4), + (16384, 768, 8), + (16384, 768, 16), + (16384, 768, 32), + (16384, 768, 64), + (16384, 768, 128), + (16384 * 8, 768, 2), + (16384 * 8, 768, 4), + (16384 * 8, 768, 8), + (16384 * 8, 768, 16), + (16384 * 8, 768, 32), + (16384 * 8, 768, 64), + (16384 * 8, 768, 128), +) + + +# class PermuteBenchmark(parameterized.TestCase): +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testBinnedGather(self, sl, hs, ne): +# # NOTE: Capacity factor == 1. +# ec = sl // ne +# +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# top_expert = torch.randint(0, ne, (sl,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(indices, ne) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# +# def benchmark(): +# return ops.binned_gather(x, indices, bins, ec) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('BinnedGather', arguments, mean_t, std_t) +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testBinnedScatter(self, sl, hs, ne): +# # NOTE: Capacity factor == 1. +# ec = sl // ne +# +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# top_expert = torch.randint(0, ne, (sl,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(indices, ne) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# x = ops.binned_gather(x, indices, bins, ec) +# +# def benchmark(): +# return ops.binned_scatter(x, indices, bins) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('BinnedScatter', arguments, mean_t, std_t) +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testPaddedGather(self, sl, hs, ne): +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# +# # Randomly assign tokens to experts. +# top_expert = torch.randint(0, ne, (sl,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(top_expert, ne) +# padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) +# padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# +# def benchmark(): +# return ops.padded_gather(x, indices, bin_ids, bins, padded_bins) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('PaddedGather', arguments, mean_t, std_t) +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testPaddedScatter(self, sl, hs, ne): +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# +# # Randomly assign tokens to experts. +# top_expert = torch.randint(0, ne, (sl,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(top_expert, ne) +# padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) +# padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins) +# +# def benchmark(): +# return ops.padded_scatter(x, indices, bin_ids, bins, padded_bins) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('PaddedScatter', arguments, mean_t, std_t) +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testCopy(self, sl, hs, ne): +# # NOTE: Capacity factor == 1. +# # ec = sl // ne +# +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# y = x.clone() +# +# def benchmark(): +# return y.copy_(x) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('Copy', arguments, mean_t, std_t) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/repeat.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/repeat.py new file mode 100644 index 0000000000000000000000000000000000000000..7e9e09de5f857d51cd758ab30b2f3a846d6f9275 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/repeat.py @@ -0,0 +1,10 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch + + +def repeat(x: torch.Tensor, tiling: torch.Size): + if all((t == 1 for t in tiling)): + return x + return x.repeat(*tiling) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/replicate.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/replicate.py new file mode 100644 index 0000000000000000000000000000000000000000..26daf0eede330603a4b8ea7167faf1411d07ca93 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/replicate.py @@ -0,0 +1,36 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + + +# Autograd wrapper for replicate kernel. +class ReplicateOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, bins: torch.Tensor, num_outputs: int): + ctx.save_for_backward(bins) + out = torch.empty((x.shape[0], num_outputs), dtype=x.dtype, device=x.device) + ops.replicate_forward(x, bins, out) + return out + + @staticmethod + def backward(ctx: Any, grad: torch.Tensor): + bins, = ctx.saved_tensors + out = torch.empty((grad.shape[0], bins.shape[0]), dtype=grad.dtype, device=grad.device) + ops.replicate_backward(grad, bins, out) + return out, None, None + + +replicate = ReplicateOp.apply diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/round_up.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/round_up.py new file mode 100644 index 0000000000000000000000000000000000000000..6cf6bc873c9f448c5fa9126ebcfd66e8688002af --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/round_up.py @@ -0,0 +1,14 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch + + +def round_up(x: torch.Tensor, value: int): + assert isinstance(value, int) + assert x.dtype == torch.int32 + + # TODO(tgale): If this becomes and issue + # do this in a custom kernel. We only expect + # to use this on arrays of less than 1k elements. + return torch.div(x + (value - 1), value, rounding_mode='trunc') * value diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/scatter.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/scatter.py new file mode 100644 index 0000000000000000000000000000000000000000..f4605d9b46f387761b070352365f223dbfe69d47 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/scatter.py @@ -0,0 +1,72 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Optional + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for scatter kernel. +class ScatterOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + top_k: int, + ) -> torch.Tensor: + maybe_x = [x] if ctx.needs_input_grad[3] else [] + ctx.save_for_backward(indices, bin_ids, weights, bins, *maybe_x) + ctx.top_k = top_k + ctx.x_shape = x.shape + return kernels.scatter(x, indices, bin_ids, weights, bins, top_k) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + saved_tensors = ctx.saved_tensors + + indices, bin_ids, weights, bins = saved_tensors[:4] + dgrad = None + if ctx.needs_input_grad[0]: + dgrad = kernels.gather( + grad, + indices, + bin_ids, + weights, + bins, + ctx.top_k, + ) + + wgrad = None + if ctx.needs_input_grad[3]: # need wgrad + x = saved_tensors[-1] + wgrad = kernels.scatter_wgrad( + x, + grad, + indices, + bin_ids, + bins, + ctx.top_k, + ) + return dgrad, None, None, wgrad, None, None, None + + +def scatter( + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + top_k: int, +) -> Optional[torch.Tensor]: + return ScatterOp.apply(x, indices, bin_ids, weights, bins, top_k) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/sort.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/sort.py new file mode 100644 index 0000000000000000000000000000000000000000..bda3bf64283e39533c2eae3627e76bb2d0262c9f --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/sort.py @@ -0,0 +1,38 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Optional, Tuple + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + +_BITS_FOR_DTYPE = { + torch.int16: 16, + torch.int32: 32, + torch.int64: 64, +} + + +# Autograd wrapper for sort kernel. +# NOTE: Does not support gradients. +class SortOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, end_bit: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]: + if end_bit is None: + end_bit = _BITS_FOR_DTYPE[x.dtype] + x_out = torch.empty_like(x) + iota_out = torch.empty_like(x) + ops.sort(x, end_bit, x_out, iota_out) + return (x_out, iota_out) + + +sort = SortOp.apply diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/sort_benchmark.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/sort_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..11043c0824c36372585f1d9f48480c2a6ef32eb6 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/sort_benchmark.py @@ -0,0 +1,85 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import numpy as np +import torch +# from absl.testing import parameterized + +from .. import ops + +_SORT_TESTS = ( + (16384, torch.int32, None), + (16384, torch.int32, 2), + (16384, torch.int32, 128), +) + +_BASELINE_SORT_TESTS = ((16384,),) + + +def numpy_dtype(dtype): + types = { + torch.int16: np.int16, + torch.int32: np.int32, + torch.int64: np.int64, + } + return types[dtype] + + +def benchmark_function(fn, iterations=10): + # Run once to get rid of startup overhead. + fn() + times = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + times = np.array(times) + return times.mean(), times.std(), times.max(), times.min() + + +def log_benchmark(arguments, mean_t, std_t): + print('=' * 60) + print('Benchmark Parameters:') + for (key, value) in arguments.items(): + print(f'{key} = {value}') + print('Results:') + print('mean / std = {:.2f}ms / {:.2f}ms'.format(mean_t, std_t)) + print('=' * 60) + + +# class SortBenchmark(parameterized.TestCase): +# +# @parameterized.parameters(*_SORT_TESTS) +# def testSort(self, n, dtype, max_val): +# if max_val is None: +# max_val = np.iinfo(numpy_dtype(dtype)).max +# end_bit = int(np.ceil(np.log2(max_val))) +# x = torch.randint(0, max_val, (n,)).cuda().to(dtype) +# +# mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.sort(x, end_bit),) +# arguments = { +# 'n': n, +# 'dtype': dtype, +# 'max_val': max_val, +# } +# log_benchmark(arguments, mean_t, std_t) +# +# @parameterized.parameters(*_BASELINE_SORT_TESTS) +# def testTorchSort(self, n): +# x = torch.randint(0, 128, (n,)).cuda().to(torch.int32) +# +# mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.sort(x)) +# arguments = { +# 'n': n, +# } +# log_benchmark(arguments, mean_t, std_t) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/stk_autocast.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/stk_autocast.py new file mode 100644 index 0000000000000000000000000000000000000000..7a3626e5e0eec51339c95a448bca84be14a2ca93 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/stk_autocast.py @@ -0,0 +1,39 @@ +# vendored from +# https://github.com/stanford-futuredata/stk/blob/736313768ef697ce13a0594a41b2512a0fbc9884/stk/backend/autocast.py +import functools +import torch + + +def _is_eligible(x): + return x.is_floating_point() and x.is_cuda and (x.dtype is not torch.float64) + + +def _cast(x, dtype): + if isinstance(x, torch.Tensor) and _is_eligible(x): + return x.to(dtype) + elif isinstance(x, map): + return {_cast(k, dtype): _cast(v, dtype) for k, v in x.items()} + elif isinstance(x, list) or isinstance(x, tuple): + return type(x)(map(lambda y: _cast(y, dtype), x)) + return x + + +def custom_fwd(fwd): + """Wrap a custom autograd function that always uses autocast dtype.""" + + @functools.wraps(fwd) + def decorate_fwd(*args, **kwargs): + if torch.is_autocast_enabled(): + with torch.autocast(device_type="cuda", enabled=False): + dtype = torch.get_autocast_gpu_dtype() + return fwd(*_cast(args, dtype), **_cast(kwargs, dtype)) + return fwd(*args, **kwargs) + return decorate_fwd + + +def custom_bwd(bwd): + @functools.wraps(bwd) + def decorate_bwd(*args, **kwargs): + with torch.autocast(device_type="cuda", enabled=False): + return bwd(*args, **kwargs) + return decorate_bwd \ No newline at end of file diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/sum.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/sum.py new file mode 100644 index 0000000000000000000000000000000000000000..e00c1aa68e1193f5b72f75a2edc37de8d505facc --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/sum.py @@ -0,0 +1,9 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +import torch + + +def sum(x: torch.Tensor, dim: int = 0): + if x.shape[dim] == 1: + return x.squeeze(dim=dim) + return x.sum(dim=dim) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/ops/topology.py b/build/torch214-cxx11-cu126-x86_64-linux/ops/topology.py new file mode 100644 index 0000000000000000000000000000000000000000..76a50d3164db20534b099dcb4d8487a7aef25d15 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/ops/topology.py @@ -0,0 +1,45 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + + +# Autograd wrapper for topology kernel. +# NOTE: Does not support gradients. +class TopologyOp(torch.autograd.Function): + + @staticmethod + def forward( + ctx: Any, + padded_bins: torch.Tensor, + block_size: int, + output_block_rows: int, + output_block_columns: int, + ): + out = torch.empty( + output_block_rows * output_block_columns, + dtype=torch.int16, + device=padded_bins.device, + ) + ops.indices( + padded_bins, + block_size, + output_block_rows, + output_block_columns, + out, + ) + return out + + +topology = TopologyOp.apply diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..73c40c267b1c3f4949e9c957a5d2c9f682dfc1a6 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/__init__.py @@ -0,0 +1,7 @@ +# import stk.random +# import stk.ops +# from stk.matrix import Matrix + +from . import random +from . import ops +from .matrix import Matrix diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/backend/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/backend/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/backend/autocast.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/backend/autocast.py new file mode 100644 index 0000000000000000000000000000000000000000..97f6e919a60f3fd579ed0215031008d14111dc96 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/backend/autocast.py @@ -0,0 +1,37 @@ +import functools +import torch + + +def _is_eligible(x): + return x.is_floating_point() and x.is_cuda and (x.dtype is not torch.float64) + + +def _cast(x, dtype): + if isinstance(x, torch.Tensor) and _is_eligible(x): + return x.to(dtype) + elif isinstance(x, map): + return {_cast(k, dtype): _cast(v, dtype) for k, v in x.items()} + elif isinstance(x, list) or isinstance(x, tuple): + return type(x)(map(lambda y: _cast(y, dtype), x)) + return x + + +def custom_fwd(fwd): + """Wrap a custom autograd function that always uses autocast dtype.""" + + @functools.wraps(fwd) + def decorate_fwd(*args, **kwargs): + if torch.is_autocast_enabled(): + with torch.autocast(device_type="cuda", enabled=False): + dtype = torch.get_autocast_gpu_dtype() + return fwd(*_cast(args, dtype), **_cast(kwargs, dtype)) + return fwd(*args, **kwargs) + return decorate_fwd + + +def custom_bwd(bwd): + @functools.wraps(bwd) + def decorate_bwd(*args, **kwargs): + with torch.autocast(device_type="cuda", enabled=False): + return bwd(*args, **kwargs) + return decorate_bwd diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/backend/sputnik.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/backend/sputnik.py new file mode 100644 index 0000000000000000000000000000000000000000..220c947bc1e932e8c77cc30f4069e9930f1aa962 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/backend/sputnik.py @@ -0,0 +1,316 @@ +import torch + +from ..backend import triton_kernels as backend +from ..backend.autocast import custom_bwd, custom_fwd + + +def _standardize_shape(x, transpose): + if transpose: + return torch.Size((x[1], x[0])) + return x + + +def _sparse_transpose(x): + return (torch.Size((x[0][1], x[0][0])), ) + x[1:] + + +def _transpose_helper(x, transpose): + if isinstance(x, torch.Tensor): + return x.t() if transpose else x + if transpose: + x = _sparse_transpose(x) + return x + (transpose,) + + +def _wrap(x): + if isinstance(x, torch.Tensor): + return (x,) + return x + + +def _is_transposed(x): + return (not x.is_contiguous() and + x.stride()[0] == 1 and + x.stride()[1] == x.size()[0]) + + +def _call_helper(op, out, a, b, trans_a, trans_b): + args = (_wrap(_transpose_helper(a, trans_a)) + + _wrap(_transpose_helper(b, trans_b))) + if isinstance(out, tuple): + args = args + out + return op(*args) + + +def _preprocess_inputs(lhs, rhs, dy): + if isinstance(lhs, torch.Tensor) and _is_transposed(lhs): + lhs = lhs.t() + if isinstance(rhs, torch.Tensor) and _is_transposed(rhs): + rhs = rhs.t() + if (isinstance(dy, torch.Tensor) and + not dy.is_contiguous() and + not _is_transposed(dy)): + dy = dy.contiguous() + if isinstance(dy, tuple) and not dy[1].is_contiguous(): + dy = (dy[0], dy[1].contiguous()) + dy[2:] + return lhs, rhs, dy + + +def _postprocess_outputs(x, transpose, grad): + if isinstance(x, torch.Tensor) and transpose: + return grad.t() + return grad + + +def _lhs_gradient(op, lhs, rhs, dy, trans_lhs, trans_rhs): + lhs, rhs, dy = _preprocess_inputs(lhs, rhs, dy) + + a, b = (rhs, dy) if trans_lhs else (dy, rhs) + trans_a = trans_lhs and trans_rhs + trans_b = trans_lhs or not trans_rhs + out = _call_helper(op, lhs, a, b, trans_a, trans_b) + return _postprocess_outputs(lhs, trans_lhs, out) + + +def _rhs_gradient(op, lhs, rhs, dy, trans_lhs, trans_rhs): + lhs, rhs, dy = _preprocess_inputs(lhs, rhs, dy) + + a, b = (dy, lhs) if trans_rhs else (lhs, dy) + trans_a = not trans_lhs or trans_rhs + trans_b = trans_lhs and trans_rhs + out = _call_helper(op, rhs, a, b, trans_a, trans_b) + return _postprocess_outputs(rhs, trans_rhs, out) + + +class DSD(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward(ctx, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_a, + rhs): + ctx.save_for_backward(data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + rhs) + ctx.shape = _standardize_shape(shape, transpose_a) + ctx.transpose_a = transpose_a + + out = torch.empty( + (shape[0], rhs.size()[1]), + dtype=rhs.dtype, + device=rhs.device) + + backend.dsd(shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_a, + rhs, + out) + return out + + @staticmethod + @custom_bwd + def backward(ctx, dy): + saved_tensors = ctx.saved_tensors + lhs = (ctx.shape,) + saved_tensors[:-1] + rhs = saved_tensors[-1] + trans_a = ctx.transpose_a + trans_b = _is_transposed(rhs) + + ddata = None + if ctx.needs_input_grad[1]: + ddata = _lhs_gradient(sdd, + lhs, + rhs, + dy, + trans_a, + trans_b) + drhs = None + if ctx.needs_input_grad[-1]: + op = dds if trans_b else dsd + drhs = _rhs_gradient(op, + lhs, + rhs, + dy, + trans_a, + trans_b) + return None, ddata, None, None, None, None, None, None, None, drhs + + +dsd = DSD.apply + + +class DDS(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward(ctx, + lhs, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_b): + ctx.save_for_backward(lhs, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t) + ctx.shape = _standardize_shape(shape, transpose_b) + ctx.transpose_b = transpose_b + out = torch.empty((lhs.size()[0], shape[1]), + dtype=lhs.dtype, + device=lhs.device) + backend.dds(lhs, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_b, + out) + return out + + @staticmethod + @custom_bwd + def backward(ctx, dy): + saved_tensors = ctx.saved_tensors + lhs = saved_tensors[0] + rhs = (ctx.shape,) + saved_tensors[1:] + trans_a = _is_transposed(lhs) + trans_b = ctx.transpose_b + + dlhs = None + if ctx.needs_input_grad[0]: + op = dsd if trans_a else dds + dlhs = _lhs_gradient(op, + lhs, + rhs, + dy, + trans_a, + trans_b) + ddata = None + if ctx.needs_input_grad[2]: + ddata = _rhs_gradient(sdd, + lhs, + rhs, + dy, + trans_a, + trans_b) + return dlhs, None, ddata, None, None, None, None, None, None, None + + +dds = DDS.apply + + +class SDD(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward(ctx, + lhs, + rhs, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t): + ctx.save_for_backward( + lhs, + rhs, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t) + ctx.shape = shape + out = torch.empty( + data.shape, + dtype=lhs.dtype, + device=lhs.device) + backend.sdd(lhs, + rhs, + shape, + out, + offsets, + row_indices, + column_indices) + return out + + @staticmethod + @custom_bwd + def backward(ctx, dy): + saved_tensors = ctx.saved_tensors + lhs, rhs = saved_tensors[:2] + dy = (ctx.shape, dy) + saved_tensors[2:] + trans_a = _is_transposed(lhs) + trans_b = _is_transposed(rhs) + + dlhs = None + if ctx.needs_input_grad[0]: + op = dds if trans_a else dsd + dlhs = _lhs_gradient(op, + lhs, + rhs, + dy, + trans_a, + trans_b) + drhs = None + if ctx.needs_input_grad[1]: + op = dsd if trans_b else dds + drhs = _rhs_gradient(op, + lhs, + rhs, + dy, + trans_a, + trans_b) + return dlhs, drhs, None, None, None, None, None, None, None, None + + +sdd = SDD.apply + +class RowIndices(torch.autograd.Function): + + @staticmethod + def forward(ctx, shape, data, offsets, column_indices): + out = torch.empty( + column_indices.shape, + dtype=column_indices.dtype, + device=column_indices.device) + backend.row_indices(shape, data, offsets, column_indices, out) + return out + + +row_indices = RowIndices.apply diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/backend/triton_kernels.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/backend/triton_kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..c535309f3321249f475367164a558f94a4f8eb86 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/backend/triton_kernels.py @@ -0,0 +1,393 @@ +import torch +import triton +import triton.language as tl +from dataclasses import dataclass + +@dataclass +class TritonConfig: + BLOCK_M: int = 128 + BLOCK_N: int = 128 + BLOCK_K: int = 32 + BLOCK_SIZE: int = 128 + NUM_STAGES: int = 4 + NUM_WARPS: int = 4 + +def _validate_matmul_dims(M: int, K: int, N: int): + error_string = "incompatible dimensions: tensor has dim with length: {}, which must be divisible by {}" + assert M % TritonConfig.BLOCK_M == 0, error_string.format(M, TritonConfig.BLOCK_M) + assert K % TritonConfig.BLOCK_K == 0, error_string.format(K, TritonConfig.BLOCK_K) + assert N % TritonConfig.BLOCK_N == 0, error_string.format(N, TritonConfig.BLOCK_N) + +@triton.autotune( + configs=[ + # basic configs for compute-bound matmuls + triton.Config({ + 'BLOCK_M': TritonConfig.BLOCK_M, + 'BLOCK_N': TritonConfig.BLOCK_N, + 'BLOCK_K': TritonConfig.BLOCK_K, + 'BLOCK_SIZE': TritonConfig.BLOCK_SIZE + }, num_stages=TritonConfig.NUM_STAGES, num_warps=TritonConfig.NUM_WARPS), + ], + key=['M', 'N', 'K'], +) +@triton.jit +def _sdd_kernel(A, B, C, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + row_indices, column_indices, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + BLOCK_SIZE: tl.constexpr, GROUP_M: tl.constexpr, ACC_TYPE: tl.constexpr, + ): + # matrix multiplication + pid = tl.program_id(0) + pid_m = tl.load(row_indices + pid) + pid_n = tl.load(column_indices + pid) + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) + rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) + rk = tl.arange(0, BLOCK_K) + # pointers + A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) + B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) + # do matrix multiplication + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) + for k in range(0, tl.cdiv(K, BLOCK_K)): + a = tl.load(A) + b = tl.load(B) + acc += tl.dot(a, b) + A += BLOCK_K * stride_ak + B += BLOCK_K * stride_bk + #Store to sparse matrix + acc = acc.to(C.dtype.element_ty) + BLOCK_ELEMENTS = BLOCK_SIZE * BLOCK_SIZE + cm = tl.arange(0, BLOCK_M) + cn = tl.arange(0, BLOCK_N) + C = C + pid * BLOCK_ELEMENTS + (cm[:, None] * stride_cm + cn[None, :] * stride_cn) + tl.store(C, acc, mask=True) + +@triton.autotune( + configs=[ + # basic configs for compute-bound matmuls + triton.Config({ + 'BLOCK_M': TritonConfig.BLOCK_M, + 'BLOCK_N': TritonConfig.BLOCK_N, + 'BLOCK_K': TritonConfig.BLOCK_K, + 'BLOCK_SIZE': TritonConfig.BLOCK_SIZE + }, num_stages=TritonConfig.NUM_STAGES, num_warps=TritonConfig.NUM_WARPS), + ], + key=['M', 'N', 'K'], +) +@triton.jit +def _dsd_kernel(A, B, C, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + row_indices, column_indices, offsets, + block_offsets_t, trans_A: tl.constexpr, trans_B: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + BLOCK_SIZE: tl.constexpr, GROUP_M: tl.constexpr, ACC_TYPE: tl.constexpr, + ): + + # matrix multiplication + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + num_pid_m = tl.num_programs(0) + num_pid_n = tl.num_programs(1) + pid_n, pid_m = tl.swizzle2d(pid_n, pid_m, num_pid_n, num_pid_m, GROUP_M) + + start_inx = tl.load(offsets + pid_m) + end_inx = tl.load(offsets + pid_m + 1) + + # pointers to sparse matrix + rm = tl.arange(0, BLOCK_M) + rak = tl.arange(0, BLOCK_K) + + A += (rm[:, None] * stride_am + rak[None, :] * stride_ak) + + # pointers to dense matrix + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + rbk = tl.arange(0, BLOCK_K) + B += (rbk[:, None] * stride_bk + rn[None, :] * stride_bn) + + # do matrix multiplication + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) + nsub_blocks = tl.cdiv(BLOCK_SIZE, BLOCK_K) + + BLOCK_ELEMENTS = BLOCK_SIZE * BLOCK_SIZE + ak_sub_incr = BLOCK_K * stride_ak + bk_sub_incr = BLOCK_K * stride_bk + bk_block_incr = BLOCK_SIZE * stride_bk + + for k in range(nsub_blocks * (end_inx - start_inx)): + sub_block_inx = k % nsub_blocks + block_inx = k // nsub_blocks + + if trans_A: + ptr_A = A + tl.load(block_offsets_t + start_inx + block_inx) * BLOCK_ELEMENTS + sub_block_inx * ak_sub_incr + else: + ptr_A = A + (start_inx + block_inx) * BLOCK_ELEMENTS + sub_block_inx * ak_sub_incr + + ptr_B = B + tl.load(column_indices + start_inx + block_inx) * bk_block_incr + sub_block_inx * bk_sub_incr + + a = tl.load(ptr_A) + b = tl.load(ptr_B) + acc += tl.dot(a, b) + + acc = acc.to(C.dtype.element_ty) + + cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + C = C + (cm[:, None] * stride_cm + cn[None, :] * stride_cn) + tl.store(C, acc, mask=True) + +@triton.autotune( + configs=[ + # basic configs for compute-bound matmuls + triton.Config({ + 'BLOCK_M': TritonConfig.BLOCK_M, + 'BLOCK_N': TritonConfig.BLOCK_N, + 'BLOCK_K': TritonConfig.BLOCK_K, + 'BLOCK_SIZE': TritonConfig.BLOCK_SIZE + }, num_stages=TritonConfig.NUM_STAGES, num_warps=TritonConfig.NUM_WARPS), + ], + key=['M', 'N', 'K'], +) +@triton.jit +def _dds_kernel(A, B, C, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + row_indices, column_indices, offsets, + block_offsets_t, trans_A: tl.constexpr, trans_B: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + BLOCK_SIZE: tl.constexpr, GROUP_M: tl.constexpr, ACC_TYPE: tl.constexpr, + ): + + # matrix multiplication + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + num_pid_m = tl.num_programs(0) + num_pid_n = tl.num_programs(1) + pid_n, pid_m = tl.swizzle2d(pid_n, pid_m, num_pid_n, num_pid_m, GROUP_M) + + start_inx = tl.load(offsets + pid_n) + end_inx = tl.load(offsets + pid_n + 1) + + # pointers to dense matrix + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rak = tl.arange(0, BLOCK_K) + + A += (rm[:, None] * stride_am + rak[None, :] * stride_ak) + + # pointers to sparse matrix + rn = tl.arange(0, BLOCK_N) + rbk = tl.arange(0, BLOCK_K) + B += (rbk[:, None] * stride_bk + rn[None, :] * stride_bn) + + # do matrix multiplication + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) + nsub_blocks = tl.cdiv(BLOCK_SIZE, BLOCK_K) + + BLOCK_ELEMENTS = BLOCK_SIZE * BLOCK_SIZE + + ak_sub_incr = BLOCK_K * stride_ak + ak_block_incr = BLOCK_SIZE * stride_ak + bk_sub_incr = BLOCK_K * stride_bk + + for k in range(nsub_blocks * (end_inx - start_inx)): + sub_block_inx = k % nsub_blocks + block_inx = k // nsub_blocks + + if trans_B: + ptr_B = B + (start_inx + block_inx) * BLOCK_ELEMENTS + sub_block_inx * bk_sub_incr + else: + ptr_B = B + tl.load(block_offsets_t + start_inx + block_inx) * BLOCK_ELEMENTS + sub_block_inx * bk_sub_incr + + ptr_A = A + tl.load(column_indices + start_inx + block_inx) * ak_block_incr + sub_block_inx * ak_sub_incr + a = tl.load(ptr_A) + b = tl.load(ptr_B) + acc += tl.dot(a, b) + + acc = acc.to(C.dtype.element_ty) + cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + C = C + (cm[:, None] * stride_cm + cn[None, :] * stride_cn) + tl.store(C, acc, mask=True) + +def dsd(shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_a, + rhs, + out + ): + + device = rhs.device + trans_A = transpose_a + trans_B = False + + if rhs.stride(0) > 1 and rhs.stride(1) > 1: + trans_B = True + + # checks constraints + assert shape[1] == rhs.shape[0], "incompatible dimensions" + M, K = shape + _, N = rhs.shape + + _validate_matmul_dims(M, K, N) + + # accumulator types + ACC_TYPE = tl.float32 if rhs.dtype in [torch.float16, torch.bfloat16, torch.float32] else tl.int32 + + stride_am, stride_ak = data.stride(1), data.stride(2) + stride_bk, stride_bn = rhs.stride(0), rhs.stride(1) + a_column_indices = column_indices + a_offsets = offsets + + # launch kernel + grid = lambda META: (triton.cdiv(M, META['BLOCK_M']), triton.cdiv(N, META['BLOCK_N'])) + + if trans_A: + stride_am, stride_ak = data.stride(2), data.stride(1) + a_column_indices, a_offsets = column_indices_t, offsets_t + + if trans_B: + stride_bk, stride_bn = rhs.stride(1), rhs.stride(0) + + _dsd_kernel[grid]( + data.data, rhs, out, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + out.stride(0), out.stride(1), + row_indices, a_column_indices, a_offsets, + block_offsets_t, trans_A, trans_B, + GROUP_M=128, ACC_TYPE=ACC_TYPE + ) + # return out + +def dds(lhs, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_b, + out + ): + + device = lhs.device + trans_B = transpose_b + trans_A = False + + if lhs.stride(0) > 1 and lhs.stride(1) > 1: + trans_A = True + + # checks constraints + assert lhs.shape[1] == shape[0], "incompatible dimensions" + M, K = lhs.shape + _, N = shape + + _validate_matmul_dims(M, K, N) + + # accumulator types + ACC_TYPE = tl.float32 if lhs.dtype in [torch.float16, torch.bfloat16, torch.float32] else tl.int32 + + stride_am, stride_ak = lhs.stride(0), lhs.stride(1) + stride_bk, stride_bn = data.stride(1), data.stride(2) + b_column_indices = column_indices_t + b_offsets = offsets_t + + # launch kernel + grid = lambda META: (triton.cdiv(M, META['BLOCK_M']), triton.cdiv(N, META['BLOCK_N'])) + + if trans_A: + stride_am, stride_ak = lhs.stride(1), lhs.stride(0) + if trans_B: + stride_bk, stride_bn = data.stride(2), data.stride(1) + b_column_indices, b_offsets = column_indices, offsets + + _dds_kernel[grid]( + lhs, data, out, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + out.stride(0), out.stride(1), + row_indices, b_column_indices, b_offsets, + block_offsets_t, trans_A, trans_B, + GROUP_M=128, ACC_TYPE=ACC_TYPE + ) + +def sdd(lhs, + rhs, + shape, + out, + offsets, + row_indices, + column_indices + ): + + device = out.device + trans_A = False + trans_B = False + + if lhs.stride(0) > 1 and lhs.stride(1) > 1: + trans_A = True + if rhs.stride(0) > 1 and rhs.stride(1) > 1: + trans_B = True + + # checks constraints + assert lhs.shape[1] == rhs.shape[0], "incompatible dimensions" + M, K = lhs.shape + _, N = rhs.shape + + _validate_matmul_dims(M, K, N) + + # accumulator types + ACC_TYPE = tl.float32 if out.dtype in [torch.float16, torch.bfloat16, torch.float32] else tl.int32 + + # launch kernel + nnz_blocks = len(row_indices) + grid = lambda META: (nnz_blocks,) + + stride_am, stride_ak = lhs.stride(0), lhs.stride(1) + stride_bk, stride_bn = rhs.stride(0), rhs.stride(1) + + if trans_A: + stride_am, stride_ak = lhs.stride(1), lhs.stride(0) + if trans_B: + stride_bk, stride_bn = rhs.stride(1), rhs.stride(0) + + _sdd_kernel[grid]( + lhs, rhs, out, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + out.stride(1), out.stride(2), + row_indices, column_indices, + GROUP_M=128, ACC_TYPE=ACC_TYPE + ) + +@triton.jit +def _row_indices_kernel(offsets, out): + pid = tl.program_id(0) + row_offset = tl.load(offsets + pid) + nnz_blocks = tl.load(offsets + pid + 1) - row_offset + for nnz_block in range(nnz_blocks): + tl.store(out + row_offset + nnz_block, pid) + +def row_indices( + shape, data, offsets, column_indices, out +): + block_rows = len(offsets) - 1 + _row_indices_kernel[(block_rows, )](offsets, out) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/matrix.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/matrix.py new file mode 100644 index 0000000000000000000000000000000000000000..80f42263d6aada287adbfa52a61fe950162a9e28 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/matrix.py @@ -0,0 +1,329 @@ +import numpy as np +import torch + +# 1. Add heavyweight (data) validation helper. +# 2. Add construction helpers +# 3. Make indentation consistent +# 4. Replace asserts with descriptive errors. + +## +### Validation helpers. +## + + +def _validate_matrix(shape, data, row_indices, column_indices, offsets): + # Data should be [nnz, block_size, block_size] + if data.dim() == 1: + data = torch.reshape(data, [data.numel(), 1, 1]) + + # Blocks should be square. + if data.shape[-2] != data.shape[-1]: + raise ValueError( + "Expected square blocking in data. " + f"Got block shape {[data.shape[-2], data.shape[-1]]}") + + # Flatten batch dimensions on data - original shape preserved + # in shape argument. + block_size = data.shape[-1] + data = data.view([-1, block_size, block_size]) + + if data.dim() != 3: + raise ValueError( + "Expected 3D shape for data (nnz, block, block). " + f"Got shape {data.dim()}D shape.") + + block_size = data.shape[1] + if shape[-2] % block_size != 0 or shape[-1] % block_size != 0: + raise ValueError( + "Matrix shape must be dividible by blocking. " + f"Got shape {shape} with " + f"{[block_size, block_size]} blocking.") + + if np.prod(shape) < data.numel(): + raise ValueError( + "Invalid matrix. Number of nonzeros exceeds matrix capacity " + f"({data.numel()} v. {np.prod(shape)})") + + if row_indices.dim() != 1: + raise ValueError( + f"Expected 1D row_indices. Got {row_indices.dim()}D row_indices.") + + if column_indices.dim() != 1: + raise ValueError( + f"Expected 1D column_indices. Got {column_indices.dim()}D column_indices.") + + if offsets.dim() != 1: + raise ValueError( + f"Expected 1D offsets. Got {offsets.dim()}D offsets.") + + if row_indices.numel() != data.shape[0]: + raise ValueError( + "Expected 1 index per nonzero block. " + f"Got {row_indices.numel()} row_indices for {data.shape[0]} blocks") + + if column_indices.numel() != data.shape[0]: + raise ValueError( + "Expected 1 index per nonzero block. " + f"Got {column_indices.numel()} column_indices for {data.shape[0]} blocks") + + block_rows = np.prod(shape[:-1]) / block_size + if offsets.numel() != block_rows + 1: + raise ValueError( + "Expected one offset per block row plus one. " + f"Got {offsets.numel()} offsets with {block_rows} block rows.") + + is_cuda = (data.is_cuda and + row_indices.is_cuda and + column_indices.is_cuda and + offsets.is_cuda) + is_cpu = (not data.is_cuda and + not row_indices.is_cuda and + not column_indices.is_cuda and + not offsets.is_cuda) + if not (is_cuda or is_cpu): + raise ValueError( + "Expected data & meta-data on common device. " + f"Got data on {data.device}, row_indices on {row_indices.device} " + f"column_indices on {column_indices.device} and " + f"offsets on {offsets.device}.") + + if data.dtype != torch.float16: + raise ValueError( + f"Expected float16 data. Got {data.dtype} data.") + if row_indices.dtype != torch.int16: + raise ValueError( + f"Expected int16 row_indices. Got {row_indices.dtype} row_indices.") + if column_indices.dtype != torch.int16: + raise ValueError( + f"Expected int16 column_indices. Got {column_indices.dtype} column_indices.") + if offsets.dtype != torch.int32: + raise ValueError( + f"Expected int32 offsets. Got {offsets.dtype} offsets.") + return data + + +def _transpose(size, data, row_indices, column_indices, offsets): + block_columns = size[1] // data.shape[1] + + # Sort row indices by column indices to get the transposed matrix's + # column indices. + gather_indices = column_indices.argsort() + column_indices_t = row_indices.gather(0, gather_indices) + block_offsets_t = gather_indices.int() + + # NOTE: Histogram is not implemented for any integer type on CPU. Do + # the histogram in 32-bit float, which can exactly represent 16-bit + # integers. + column_indices_float = column_indices.float() + + zero = torch.zeros((1,), dtype=torch.int32, device=data.device) + nnz_per_column = column_indices_float.histc(block_columns, 0, block_columns) + nnz_per_column = nnz_per_column.int() + offsets_t = torch.cat([zero, nnz_per_column.cumsum(0, dtype=torch.int32)]) + return column_indices_t, offsets_t, block_offsets_t + + +class Matrix(torch.nn.Module): + """A matrix stored in sparse format. + + Underlying format is block compressed sparse row (BCSR). + + TODO(tgale): Make this mirror torch.Tensor API as much as possible. + """ + + def __init__(self, + size, + data, + row_indices, + column_indices, + offsets, + column_indices_t=None, + offsets_t=None, + block_offsets_t=None): + super().__init__() + self._size = size + self._data = data + self._row_indices = row_indices + self._column_indices = column_indices + self._offsets = offsets + + # Produce the transpose meta-data if it is not passed in. + if ((column_indices_t is None) or (offsets_t is None) or + (block_offsets_t is None)): + column_indices_t, offsets_t, block_offsets_t = _transpose( + size, data, row_indices, column_indices, offsets) + self._column_indices_t = column_indices_t + self._offsets_t = offsets_t + self._block_offsets_t = block_offsets_t + + self._transposed = False + + # Validate that our metadata will not overflow. + max_dim = np.iinfo(np.int16).max * self.blocking + if column_indices.dtype == torch.int16: + if size[0] > max_dim or size[1] > max_dim: + raise ValueError( + "Sparse matrix with shape {size} exceeds representable " + "size with 16-bit indices.") + + def validate(self): + _validate_matrix(self._size, + self._data, + self._row_indices, + self._column_indices, + self._offsets) + + # TODO(tgale): Add heavyweight data validation. + + def to(self, device): + # TODO(tgale): Handle type conversions here. We + # need to set the appropriate meta-data type for + # the given floating-point type. + self._data = self._data.to(device) + self._row_indices = self._row_indices.to(device) + self._column_indices = self._column_indices.to(device) + self._offsets = self._offsets.to(device) + self._column_indices_t = self._column_indices_t.to(device) + self._offsets_t = self._offsets_t.to(device) + self._block_offsets_t = self._block_offsets_t.to(device) + return self + + def cuda(self): + return self.to(torch.cuda.current_device()) + + def clone(self): + return Matrix( + self.size(), + self.data.clone(), + self.row_indices.clone(), + self.column_indices.clone(), + self.offsets.clone(), + self.column_indices_t.clone(), + self.offsets_t.clone(), + self.block_offsets_t.clone()) + + def t(self): + if self.dim() != 2: + raise ValueError( + "t() expects a tensor with <= 2 dimensions, " + f"but self is {self.dim()}D.") + out = Matrix(self.size(), + self.data, + self.row_indices, + self.column_indices, + self.offsets, + self.column_indices_t, + self.offsets_t, + self.block_offsets_t) + out._transposed = not self._transposed + out._size = torch.Size((self._size[1], self._size[0])) + return out + + def contiguous(self): + raise ValueError("Not yet implemented.") + + def is_contiguous(self): + return not self._transposed + + @property + def is_cuda(self): + return self._data.is_cuda + + @property + def device(self): + return self._data.device + + def size(self): + return self._size + + @property + def shape(self): + return self.size() + + def dim(self): + return len(self._size) + + @property + def data(self): + return self._data + + @property + def row_indices(self): + return self._row_indices + + @property + def column_indices(self): + return self._column_indices + + @property + def offsets(self): + return self._offsets + + @property + def offsets_t(self): + return self._offsets_t + + @property + def column_indices_t(self): + return self._column_indices_t + + @property + def block_offsets_t(self): + return self._block_offsets_t + + @property + def dtype(self): + return self.data.dtype + + @property + def nnz(self): + return self.data.numel() + + @property + def blocking(self): + return self.data.shape[1] + + @property + def requires_grad(self): + return self.data.requires_grad + + def requires_grad_(self, x): + self.data.requires_grad_(x) + return self + + def view(self, *shape): + assert self.is_contiguous() + if shape[-1] != self.size()[-1]: + raise ValueError( + "Can't change view on compressed dimension. " + f"{self.size()[-1]} v. {shape[-1]}.") + if np.prod(shape) != np.prod(self.size()): + raise ValueError( + "Mismatch in numel of Matrix and new shape. " + f"{np.prod(self.size())} v. {np.prod(shape)}") + return Matrix(shape, + self.data, + self.row_indices, + self.column_indices, + self.offsets, + self.column_indices_t, + self.offsets_t, + self.block_offsets_t) + + @property + def grad(self): + # TODO(tgale): Make sure this mirrors torch.Tensor + # behavior in the case where we ask for the gradient + # of a non-contiguous tensor. + size = self.size() + if not self.is_contiguous(): + size = torch.Size((size[1], size[0])) + out = Matrix(size, + self.data.grad, + self.row_indices, + self.column_indices, + self.offsets, + self.column_indices_t, + self.offsets_t, + self.block_offsets_t) + return out if self.is_contiguous() else out.t() diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fc873b236f4cd4036964c016a4036e3ce5ebf1ac --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/__init__.py @@ -0,0 +1,3 @@ +from .linear_ops import dds, dsd, sdd +from .matrix_ops import ones_like, row_indices, sum, to_dense, to_sparse +from .eltwise_ops import mul diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/eltwise_ops.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/eltwise_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ba7d7332320250fd01fa60e60528f19de3e8ed03 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/eltwise_ops.py @@ -0,0 +1,28 @@ +from ..matrix import Matrix + +def mul(a, b): + """Performs element-wise multiplication of matrices a and b. + + It is the user's responsibility to make sure that a and b + follow the same matrix topology. This function assumes it is safe + to use the topoplogy of a. + + Args: + a: stk.Matrix. + b: stk.Matrix with a's matrix topology. + + Returns: + stk.Matrix where the entries correspond to torch.mul(a, b). + """ + assert isinstance(a, Matrix) + assert isinstance(b, Matrix) + assert a.size() == b.size() + + return Matrix(a.size(), + a.data * b.data, + a.row_indices, + a.column_indices, + a.offsets, + a.column_indices_t, + a.offsets_t, + a.block_offsets_t) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/eltwise_ops_test.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/eltwise_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2939372a5c68ac92b47b11015db4f75f4fd60ffa --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/eltwise_ops_test.py @@ -0,0 +1,86 @@ +import unittest +import itertools +import torch +# from absl.testing import parameterized + +import stk +from stk.ops.linear_ops_test import allclose, _dense_and_sparse + +_MATRIX_SIZES = ( + (128, 128, 0.0), + (256, 256, 0.5), + (2048, 1024, 0.8), + (512, 128, 0.0), + (128, 512, 0.0), + (1024, 512, 0.0), + (1024, 512, 0.5), + (1024, 512, 0.75), + (512, 1024, 0.0), + (512, 1024, 0.5), + (512, 1024, 0.75), + (1024, 1024, 0.0), + (1024, 1024, 0.5), + (1024, 1024, 0.75), +) + +_DTYPE = ( + torch.float16, torch.bfloat16 +) + +def _generate_testcases(): + testcases = itertools.product(_MATRIX_SIZES, _DTYPE) + testcases = [(*size, 128, dtype) for + (size, dtype) in testcases] + return testcases + +_ELTWISE_OP_TESTS = _generate_testcases() + +def _dense_and_sparse_like(x, std=0.1): + dense_data = torch.randn_like(x.data, device=x.device) * std + sparse = stk.Matrix(x.size(), + dense_data, + x.row_indices, + x.column_indices, + x.offsets) + dense = stk.ops.to_dense(sparse) + + return (dense.requires_grad_(True), + sparse.requires_grad_(True)) + +# @parameterized.parameters(_ELTWISE_OP_TESTS) +# class EltwiseOpsTest(parameterized.TestCase): +# +# def testEltwiseMul(self, m, n, sparsity, blocking, dtype): +# +# a_dense, a = _dense_and_sparse(m, n, sparsity, blocking, dtype) +# b_dense, b = _dense_and_sparse_like(a) +# +# out = stk.ops.mul(a, b) +# expected_out = torch.mul(a_dense, b_dense) +# +# # Compute the gradients w.r.t. the inputs. +# expected_out.sum().backward() +# stk.ops.sum(out).backward() +# +# # Validate the results. +# out = stk.ops.to_dense(out) +# self.assertEqual(out.dim(), 2) +# self.assertEqual(expected_out.size(), out.size()) +# self.assertTrue(allclose(out, expected_out)) +# +# # LHS gradient. +# grad = stk.ops.to_dense(a.grad) +# expected_grad = a_dense.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size(), grad.size()) +# self.assertTrue(allclose(grad, expected_grad)) +# +# # RHS gradient. +# grad = stk.ops.to_dense(b.grad) +# expected_grad = b_dense.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size(), grad.size()) +# self.assertTrue(allclose(grad, expected_grad)) + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/linear_ops.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/linear_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..9d277c8c07f9e30addc31900a12175c8a1f4d7ad --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/linear_ops.py @@ -0,0 +1,59 @@ +import torch + +from ..backend import sputnik +from ..matrix import Matrix + + +def dsd(a, b): + assert isinstance(a, Matrix) + assert isinstance(b, torch.Tensor) + return sputnik.dsd( + a.size(), + a.data, a.offsets, + a.row_indices, + a.column_indices, + a.offsets_t, + a.column_indices_t, + a.block_offsets_t, + not a.is_contiguous(), + b) + + +def dds(a, b): + assert isinstance(a, torch.Tensor) + assert isinstance(b, Matrix) + return sputnik.dds( + a, + b.size(), + b.data, b.offsets, + b.row_indices, + b.column_indices, + b.offsets_t, + b.column_indices_t, + b.block_offsets_t, + not b.is_contiguous()) + + +def sdd(a, b, topo): + assert isinstance(a, torch.Tensor) + assert isinstance(b, torch.Tensor) + assert isinstance(topo, Matrix) + assert topo.is_contiguous() + out = sputnik.sdd( + a, b, + topo.size(), + topo.data, + topo.offsets, + topo.row_indices, + topo.column_indices, + topo.offsets_t, + topo.column_indices_t, + topo.block_offsets_t) + return Matrix(topo.size(), + out, + topo.row_indices, + topo.column_indices, + topo.offsets, + topo.column_indices_t, + topo.offsets_t, + topo.block_offsets_t) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/linear_ops_test.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/linear_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..c1c24d350df9c1b2346c7da885502cd696c88867 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/linear_ops_test.py @@ -0,0 +1,216 @@ +import unittest +import itertools +import numpy as np +import torch +# from absl.testing import parameterized + +import stk + + +def allclose(x, y, pct=0.25): + mask = torch.isclose(x, y, rtol=5e-2) + pct_diff = (mask.numel() - mask.sum()) / mask.numel() * 100 + if pct_diff > pct: + print("{:.2f}% of values not close.".format(pct_diff)) + return False + return True + + +# An assortment of problems designed to make sure +# the bindings are operating correctly. +_MATRIX_SIZES = ( + (128, 128, 128, 0.0), + (256, 256, 256, 0.5), + (2048, 1024, 512, 0.8), + (512, 128, 128, 0.0), + (128, 128, 512, 0.0), + (1024, 512, 512, 0.0), + (1024, 512, 512, 0.5), + (1024, 512, 512, 0.75), + (512, 512, 1024, 0.0), + (512, 512, 1024, 0.5), + (512, 512, 1024, 0.75), + (1024, 1024, 1024, 0.0), + (1024, 1024, 1024, 0.5), + (1024, 1024, 1024, 0.75), +) + +_TRANSPOSE = ( + (False, False), + (False, True), + (True, False), + (True, True), +) + +_DTYPE = ( + torch.float16, torch.bfloat16 +) + +def _generate_testcases(): + testcases = itertools.product(_MATRIX_SIZES, _TRANSPOSE, _DTYPE) + testcases = [(*size, *trans, 128, dtype) for + (size, trans, dtype) in testcases] + return testcases + +_LINEAR_OP_TESTS = _generate_testcases() + +def _dense_and_sparse(rows, cols, sparsity, blocking, dtype, std=0.1): + mask = stk.random.dense_mask(rows, cols, sparsity, blocking) + dense = (torch.randn(rows, cols) * std * mask).type(dtype) + sparse = stk.ops.to_sparse(dense, blocking) + cuda_device = torch.device("cuda") + return (dense.to(cuda_device).requires_grad_(True), + sparse.to(cuda_device).requires_grad_(True)) + + +def _dense(rows, cols, dtype, std=0.1): + cuda_device = torch.device("cuda") + out = (torch.randn(rows, cols) * std).type(dtype) + return out.to(cuda_device).requires_grad_(True) + + +def _dense_2x(rows, cols, dtype): + a = _dense(rows, cols, dtype) + return a, a.detach().requires_grad_(True) + + +def _with_transpose(op, a, b, trans_a, trans_b): + a = a.t() if trans_a else a + b = b.t() if trans_b else b + return op(a, b) + + +def _mmm(a, b, topo): + mask = stk.ops.to_dense(stk.ops.ones_like(topo)) + return torch.mm(a, b) * mask + + +def _sparse_out_with_transpose(op, a, b, topo, trans_a, trans_b): + a = a.t() if trans_a else a + b = b.t() if trans_b else b + return op(a, b, topo) + + +def _mask(x, mask): + mask = stk.ops.to_dense(stk.ops.ones_like(mask)) + return x * mask + + +# @parameterized.parameters(*_LINEAR_OP_TESTS) +# class LinearOpsTest(parameterized.TestCase): +# +# def testLinearOps_Dsd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype): +# # Construct the operands. +# a_shape = (k, m) if trans_a else (m, k) +# a_dense, a = _dense_and_sparse(*a_shape, sparsity, blocking, dtype) +# b_shape = (n, k) if trans_b else (k, n) +# b, bcp = _dense_2x(*b_shape, dtype) +# +# # Execute the matmul. +# out = _with_transpose(stk.ops.dsd, a, b, trans_a, trans_b) +# expected_out = _with_transpose(torch.mm, a_dense, bcp, trans_a, trans_b) +# +# # Compute the gradients w.r.t. the inputs. +# expected_out.sum().backward() +# out.sum().backward() +# +# # Validate the results. +# self.assertEqual(out.dim(), 2) +# self.assertEqual(expected_out.size()[0], out.size()[0]) +# self.assertEqual(expected_out.size()[1], out.size()[1]) +# self.assertTrue(allclose(out, expected_out)) +# +# # LHS gradient. +# grad = stk.ops.to_dense(a.grad) +# expected_grad = _mask(a_dense.grad, a.grad) +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# # RHS gradient. +# grad = b.grad +# expected_grad = bcp.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# def testLinearOps_Dds(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype): +# # Construct the operands. +# a_shape = (k, m) if trans_a else (m, k) +# a, acp = _dense_2x(*a_shape, dtype) +# b_shape = (n, k) if trans_b else (k, n) +# b_dense, b = _dense_and_sparse(*b_shape, sparsity, blocking, dtype) +# +# # Execute the matmul. +# out = _with_transpose(stk.ops.dds, a, b, trans_a, trans_b) +# expected_out = _with_transpose(torch.mm, acp, b_dense, trans_a, trans_b) +# +# # Compute the gradients w.r.t. the inputs. +# expected_out.sum().backward() +# out.sum().backward() +# +# # Validate the results. +# self.assertEqual(out.dim(), 2) +# self.assertEqual(expected_out.size()[0], out.size()[0]) +# self.assertEqual(expected_out.size()[1], out.size()[1]) +# self.assertTrue(allclose(out, expected_out)) +# +# # LHS gradient. +# grad = a.grad +# expected_grad = acp.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# # RHS gradient. +# grad = stk.ops.to_dense(b.grad) +# expected_grad = _mask(b_dense.grad, b.grad) +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# def testLinearOps_Sdd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype): +# # Construct the operands. +# a_shape = (k, m) if trans_a else (m, k) +# a, acp = _dense_2x(*a_shape, dtype) +# b_shape = (n, k) if trans_b else (k, n) +# b, bcp = _dense_2x(*b_shape, dtype) +# _, topo = _dense_and_sparse(m, n, sparsity, blocking, dtype) +# +# # Execute the matmul. +# out = _sparse_out_with_transpose(stk.ops.sdd, a, b, topo, trans_a, trans_b) +# expected_out = _sparse_out_with_transpose(_mmm, acp, bcp, topo, trans_a, trans_b) +# +# # Compute the gradients w.r.t. the inputs. +# expected_out.sum().backward() +# stk.ops.sum(out).backward() +# +# # Validate the results. +# out = stk.ops.to_dense(out) +# self.assertEqual(out.dim(), 2) +# self.assertEqual(expected_out.size()[0], out.size()[0]) +# self.assertEqual(expected_out.size()[1], out.size()[1]) +# self.assertTrue(allclose(out, expected_out)) +# +# # LHS gradient. +# grad = a.grad +# expected_grad = acp.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# # RHS gradient. +# grad = b.grad +# expected_grad = bcp.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/matrix_ops.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/matrix_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..447c72dc73439d84f58c917676cc04e64f13e97d --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/matrix_ops.py @@ -0,0 +1,98 @@ +from ..backend import sputnik +from ..matrix import Matrix +import torch +import numpy as np + + +@torch.no_grad() +def row_indices(shape, data, offsets, column_indices): + return sputnik.row_indices(shape, data, offsets, column_indices) + + +# TODO(tgale): Replace this helper with a custom kernel. This operation +# is much simpler to do than how it's currently implemented. +@torch.no_grad() +def _expand_for_blocking(idxs, blocking): + # Duplicate for block column dimension. + idxs = torch.reshape(idxs, [idxs.size()[0], 1, 2]).repeat(1, blocking, 1) + + # Update the column indices. + idxs[:, :, 1] *= blocking + idxs[:, :, 1] += torch.reshape(torch.arange(blocking, device=idxs.device), [1, blocking]) + + # Duplicate for block row dimension. + idxs = torch.reshape(idxs, [idxs.size()[0], 1, blocking, 2]) + idxs = idxs.repeat(1, blocking, 1, 1) + + # Update the row indices. + idxs[:, :, :, 0] *= blocking + idxs[:, :, :, 0] += torch.reshape(torch.arange(blocking, device=idxs.device), [1, blocking, 1]) + idxs = torch.reshape(idxs, [-1, 2]) + return idxs + + +# TODO(tgale): Add input type checking. +@torch.no_grad() +def to_dense(x): + assert isinstance(x, Matrix) + + shape = (np.prod(x.shape[:-1]), x.shape[-1]) + row_idxs = x.row_indices.type(torch.int32) + col_idxs = x.column_indices.type(torch.int32) + indices = _expand_for_blocking(torch.stack([row_idxs, col_idxs], dim=1), x.blocking) + indices = (indices[:, 0] * shape[1] + indices[:, 1]).type(torch.int64) + + out = torch.zeros(shape[0] * shape[1], dtype=x.dtype, device=x.device) + out.scatter_(0, indices, x.data.flatten()) + return out.reshape(x.size()) + + +@torch.no_grad() +def _mask(x, blocking=1): + assert x.dim() == 2 + assert x.size()[0] % blocking == 0 + assert x.size()[1] % blocking == 0 + block_rows = x.size()[0] // blocking + block_cols = x.size()[1] // blocking + x = torch.reshape(x, [block_rows, blocking, block_cols, blocking]) + x = torch.sum(torch.abs(x), dim=(1, 3)) + return x != 0 + + +# TODO(tgale): Add input type checking. +@torch.no_grad() +def to_sparse(x, blocking=1): + m = _mask(x, blocking) + + # TODO(tgale): Set to appropriate type for input matrix. + row_nnzs = torch.sum(m, dim=1).type(torch.int32) + zeros = torch.zeros((1,), dtype=row_nnzs.dtype, device=row_nnzs.device) + offsets = torch.cat([zeros, torch.cumsum(row_nnzs, dim=0)]) + offsets = offsets.type(torch.int32) + + indices = torch.nonzero(m).type(torch.int16) + row_indices = indices[:, 0] + column_indices = indices[:, 1] + + # Nonzero indices in the dense matrix. + nonzero_indices = torch.nonzero(m) + nonzero_indices = _expand_for_blocking(nonzero_indices, blocking) + nonzero_indices = nonzero_indices[:, 0] * x.size()[1] + nonzero_indices[:, 1] + + # Gather the data and construct the sparse matrix. + data = torch.gather(x.flatten(), dim=0, index=nonzero_indices) + data = torch.reshape(data, [-1, blocking, blocking]) + return Matrix(x.size(), data, row_indices, column_indices, offsets) + + +@torch.no_grad() +def ones_like(x): + return Matrix(x.size(), + torch.ones_like(x.data), + x.row_indices, + x.column_indices, x.offsets) + + +def sum(x): + assert isinstance(x, Matrix) + return x.data.sum() diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/matrix_ops_test.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/matrix_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..d172d921f6f08b0e4fb709207a458b0e1e071bd0 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/ops/matrix_ops_test.py @@ -0,0 +1,62 @@ +import unittest + +# from absl.testing import parameterized +import stk +import torch + + +# @parameterized.parameters( +# (8, 16, 0.0, 1), +# (8, 16, 0.5, 1), +# (8, 16, .95, 1), +# (16, 8, 0.0, 1), +# (16, 8, 0.5, 1), +# (16, 8, .95, 1), +# (8, 16, 0.0, 8), +# (8, 16, 0.5, 8), +# (8, 16, 1.0, 8), +# (16, 8, 0.0, 8), +# (16, 8, 0.5, 8), +# (16, 8, 1.0, 8), +# (128, 256, 0.5, 16), +# (256, 128, 0.75, 32), +# (512, 512, .875, 128)) +# class MatrixOpsTest(parameterized.TestCase): +# +# def testMatrixOps_FormatConversion(self, rows, cols, sparsity, blocking): +# mask = stk.random.dense_mask(rows, cols, sparsity, blocking) +# x = (torch.randn(rows, cols) * mask).type(torch.float16) +# +# # Convert the matrix to sparse format. +# sparse_x = stk.ops.to_sparse(x, blocking) +# +# # Validate the matrix. +# sparse_x.validate() +# +# # Validate the shape. +# self.assertEqual(sparse_x.dim(), 2) +# self.assertEqual(sparse_x.size()[0], rows) +# self.assertEqual(sparse_x.size()[1], cols) +# +# # Validate the sparsity. +# numblocks = rows // blocking * cols // blocking +# nnz = round(numblocks * (1 - sparsity)) * blocking ** 2 +# self.assertEqual(sparse_x.nnz, nnz) +# +# # Convert back to dense format. +# dense_x = stk.ops.to_dense(sparse_x) +# +# # Validate the shape. +# self.assertEqual(dense_x.dim(), 2) +# self.assertEqual(dense_x.size()[0], rows) +# self.assertEqual(dense_x.size()[1], cols) +# +# # Validate the sparsity +# self.assertEqual(torch.count_nonzero(dense_x).item(), nnz) +# +# # Validate the output. +# self.assertTrue(torch.all(torch.eq(x, dense_x))) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/random/__init__.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/random/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2576d1ca27283f77569a9a620c7c99fa68aaf30e --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/random/__init__.py @@ -0,0 +1,2 @@ +# from stk.random.random_ops import dense_mask, mask, randn +from .random_ops import dense_mask, mask, randn diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/random/random_ops.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/random/random_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..d1b36771e0eb8e7abf46bcb3b136b5fb1d29df93 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/random/random_ops.py @@ -0,0 +1,36 @@ +import numpy as np +import torch +from ..ops import matrix_ops + + +@torch.no_grad() +def dense_mask(rows, cols, sparsity, blocking=1): + assert sparsity >= 0.0 and sparsity <= 1.0 + assert rows % blocking == 0 and cols % blocking == 0 + + block_rows, block_cols = (rows // blocking, cols // blocking) + nnz = round(block_rows * block_cols * (1 - sparsity)) + + out = np.ones(block_rows * block_cols) + mask = np.random.choice(out.size, out.size - nnz, replace=False) + out[mask] = 0.0 + + out = np.tile( + np.reshape(out, [block_rows, 1, block_cols, 1]), + (1, blocking, 1, blocking)) + out = np.reshape(out, [rows, cols]) + return torch.from_numpy(out.astype(np.float32)) + + +@torch.no_grad() +def mask(m, n, sparsity, blocking=1): + out = dense_mask(m, n, sparsity, blocking).type(torch.float16) + return matrix_ops.to_sparse(out, blocking=blocking) + + +@torch.no_grad() +def randn(shape, sparsity, blocking=1): + shape_2d = (np.prod(shape[:-1]), shape[-1]) + out = mask(*shape_2d, sparsity, blocking) + out.data.copy_(torch.randn(*out.data.shape)) + return out.view(*shape) diff --git a/build/torch214-cxx11-cu126-x86_64-linux/stk/random/random_ops_test.py b/build/torch214-cxx11-cu126-x86_64-linux/stk/random/random_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..71d716b78b5ec009cbf9ac2dfdf09162a0102e62 --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/stk/random/random_ops_test.py @@ -0,0 +1,73 @@ +import unittest + +# from absl.testing import parameterized +from . import random +import torch + + +# @parameterized.parameters( +# (8, 16, 0.0, 1), +# (8, 16, 0.5, 1), +# (8, 16, .95, 1), +# (16, 8, 0.0, 1), +# (16, 8, 0.5, 1), +# (16, 8, .95, 1), +# (8, 16, 0.0, 8), +# (8, 16, 0.5, 8), +# (8, 16, 1.0, 8), +# (16, 8, 0.0, 8), +# (16, 8, 0.5, 8), +# (16, 8, 1.0, 8), +# (128, 256, 0.5, 16), +# (256, 128, 0.75, 32), +# (512, 512, .875, 128)) +# class RandomOpsTest(parameterized.TestCase): +# +# def testRandomOps_DenseMask(self, rows, cols, sparsity, blocking): +# mask = random.dense_mask( +# rows, cols, sparsity, blocking) +# +# # Validate the shape. +# self.assertEqual(mask.dim(), 2) +# self.assertEqual(mask.size()[0], rows) +# self.assertEqual(mask.size()[1], cols) +# +# # Validate the sparsity +# numblocks = rows // blocking * cols // blocking +# nnz = round(numblocks * (1 - sparsity)) * blocking ** 2 +# self.assertEqual( +# torch.count_nonzero(mask).item(), +# nnz) +# +# # Check values are zero or one. +# self.assertTrue( +# torch.all(torch.logical_or( +# torch.eq(mask, 0), +# torch.eq(mask, 1)))) +# +# def testRandomOps_SparseMask(self, rows, cols, sparsity, blocking): +# mask = random.mask( +# rows, cols, sparsity, blocking) +# +# # Validate the matrix. +# mask.validate() +# +# # Validate the shape. +# self.assertEqual(mask.dim(), 2) +# self.assertEqual(mask.size()[0], rows) +# self.assertEqual(mask.size()[1], cols) +# +# # Validate the sparsity. +# numblocks = rows // blocking * cols // blocking +# nnz = round(numblocks * (1 - sparsity)) * blocking ** 2 +# self.assertEqual(mask.nnz, nnz) +# +# # Check values are zero or one. +# self.assertTrue( +# torch.all(torch.logical_or( +# torch.eq(mask.data, 0), +# torch.eq(mask.data, 1)))) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu126-x86_64-linux/xpu_fused_moe.py b/build/torch214-cxx11-cu126-x86_64-linux/xpu_fused_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..b709e394e2e0461d858dec0f58176c0543a15e3c --- /dev/null +++ b/build/torch214-cxx11-cu126-x86_64-linux/xpu_fused_moe.py @@ -0,0 +1,717 @@ +# SPDX-License-Identifier: Apache-2.0 +# MegaBlocks XPU Fused MoE Implementation +import os +import torch + +from ._ops import ops, add_op_namespace_prefix + +from torch.library import register_fake + + +def resolve_dtensor(weight: torch.Tensor): + """Convert DTensor to local tensor for use with custom ops.""" + from torch.distributed._tensor import DTensor + if isinstance(weight, DTensor): + return weight.to_local() + return weight + + +# Register fake/meta kernels for torch.compile compatibility +def _register_xpu_fake_kernels(): + """Register fake kernels for XPU MoE operations to support torch.compile.""" + + def _register_if_available(op_name, fn): + if hasattr(ops, op_name): + register_fake(add_op_namespace_prefix(op_name))(fn) + + _register_if_available( + "cutlass_grouped_gemm_interface", + lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4, is_B_mxfp8=False: ptr_D, + ) + + _register_if_available( + "fused_moe_prologue", + lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, ep_rank, ep_size, num_experts_on_rank: None, + ) + + _register_if_available( + "moe_gather", + lambda output, moe_output, topk_weights, permuted_row_to_unpermuted_row, unpermuted_row_to_permuted_row, expert_first_token_offset, num_experts: None, + ) + + _register_if_available( + "silu_and_mul", + lambda out, input: None, + ) + _register_if_available( + "mul_and_silu", + lambda out, input: None, + ) + _register_if_available( + "gelu_and_mul", + lambda out, input: None, + ) + _register_if_available( + "gelu_tanh_and_mul", + lambda out, input: None, + ) + _register_if_available( + "gelu_fast", + lambda out, input: None, + ) + _register_if_available( + "gelu_new", + lambda out, input: None, + ) + _register_if_available( + "gelu_quick", + lambda out, input: None, + ) + _register_if_available( + "swigluoai_and_mul", + lambda out, input, alpha=1.702, limit=7.0: None, + ) + + +# Register fake kernels on module load +_register_xpu_fake_kernels() + + +# default +def cutlass_grouped_gemm(input_A, input_B, bias, output, expert_token_count, n, + k, num_experts): + # expert_token_count_ = torch.tensor(expert_token_count, + # dtype=torch.int64, + # device=input_A.device) + # if bias is not None: + # bias = bias.repeat_interleave(expert_token_count_, dim=0).float() + + def exclusive_prefix_sum(arr): + prefix = [0] + for i, x in enumerate(arr): + prefix.append(prefix[-1] + x) + return prefix + + expert_offset = torch.tensor(exclusive_prefix_sum(expert_token_count), + dtype=torch.int64, + device="xpu") + ops.cutlass_grouped_gemm_interface( + ptr_A=input_A, + ptr_B=input_B, + ptr_scales=None, + ptr_bias=bias, + ptr_D=output, + expert_first_token_offset=expert_offset, + N=n, + K=k, + num_experts=num_experts, + is_B_int4=False, + is_B_mxfp4=False, + is_B_mxfp8=False) + + +def cutlass_grouped_gemm_xe2(input_A, input_B, scales, bias, output, + num_rows_per_expert, n, k, num_experts, is_B_int4, + is_B_mxfp4, is_B_mxfp8=False): + expert_first_token_offset = torch.cat([ + torch.tensor([0], + dtype=num_rows_per_expert.dtype, + device=num_rows_per_expert.device), + torch.cumsum(num_rows_per_expert, dim=0) + ]).to(torch.int64) + ops.cutlass_grouped_gemm_interface( + ptr_A=input_A, + ptr_B=input_B, + ptr_scales=scales, + ptr_bias=bias, + ptr_D=output, + expert_first_token_offset=expert_first_token_offset, + N=n, + K=k, + num_experts=num_experts, + is_B_int4=is_B_int4, + is_B_mxfp4=is_B_mxfp4, + is_B_mxfp8=is_B_mxfp8) + + +def ceilDiv(a, b): + return (a + b - 1) // b + + +def compute_num_tokens_per_block(num_tokens, num_experts_per_node): + for num_tokens_per_block in [32, 64, 128, 256, 512, 1024]: + num_blocks_per_seq = ceilDiv(num_tokens, num_tokens_per_block) + if num_blocks_per_seq * num_experts_per_node <= num_tokens_per_block: + return num_tokens_per_block + return 1024 + + +def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """Reinterpret a uint8 buffer as a typed tensor by copying bytes. + + This avoids `Tensor.view(dtype)` which can fail under torch.compile + constant folding when shape divisibility is not proven. + """ + if byte_tensor.dtype != torch.uint8: + raise ValueError("byte_tensor must be uint8") + itemsize = torch.empty((), dtype=dtype).element_size() + numel = byte_tensor.numel() // itemsize + out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device) + out.view(torch.uint8).copy_(byte_tensor.contiguous()) + return out + + +def implement_zp(qweight): + # change u4 to s4 to avoid zero point in gemm kernel + # only support default zero point now + assert qweight.dtype == torch.uint8, "Input tensor must be uint8" + + high_u4 = (qweight >> 4) & 0x0F + low_u4 = qweight & 0x0F + + high_s8 = high_u4.to(torch.int8) + low_s8 = low_u4.to(torch.int8) + + high_s8 = high_s8 - 8 + low_s8 = low_s8 - 8 + + def pack_compact(a, b): + + def process_number(x): + sign = (x < 0).to(torch.uint8) + abs_low3 = (x.view(torch.uint8) & 0x7).to(torch.uint8) + return (sign << 3) | abs_low3 + + packed_a = process_number(a) + packed_b = process_number(b) + + return (packed_a << 4) | packed_b + + result = pack_compact(high_s8, low_s8) + + return result + + +def xpu_fused_moe(hidden_states, + w13, + w13_scales, + w13_bias, + w2, + w2_scales, + w2_bias, + topk_weights, + topk_ids, + n_experts_per_token, + activation, + num_experts, + ep_rank=0, + ep_size=1, + is_fp8=False, + is_int4=False, + is_mxfp4=False, + is_mxfp8=False): + ''' + hidden_states: [num_rows, hidden_size] + w13: [num_experts, 2*inter_size, hidden_size] + w13_scales: + None for bf16/fp16 + or [num_experts] for fp8 + or [num_experts, 2*inter_size, hidden_size // group_size] for 4bits + w13_bias: [num_experts, 2*inter_size] or None + w2: [num_experts, hidden_size, inter_size] + w2_scales: + None for bf16/fp16 + or [num_experts] for fp8 + or [num_experts, hidden_size, inter_size // group_size] for 4bits + w2_bias: [num_experts, hidden_size] or None + topk_weights: [num_rows, topk] + topk_ids: [num_rows, topk] + n_experts_per_token: int + activation: str + num_experts: int + is_int4: bool + is_mxfp4: bool + is_mxfp8: bool + ''' + + # Resolve DTensors to local tensors before passing to custom ops + hidden_states = resolve_dtensor(hidden_states) + w13 = resolve_dtensor(w13) + w2 = resolve_dtensor(w2) + if w13_scales is not None: + w13_scales = resolve_dtensor(w13_scales) + if w13_bias is not None: + w13_bias = resolve_dtensor(w13_bias) + if w2_scales is not None: + w2_scales = resolve_dtensor(w2_scales) + if w2_bias is not None: + w2_bias = resolve_dtensor(w2_bias) + topk_weights = resolve_dtensor(topk_weights) + topk_ids = resolve_dtensor(topk_ids) + + output = torch.empty_like(hidden_states) + num_rows, hidden_size = list(hidden_states.shape) + + dim_last = w13.shape[-1] + dim_second_last = w13.shape[-2] + + # w13 is combined gate+up weights, so one dimension is 2*inter_size + # Determine which dimension is hidden_size and which is 2*inter_size + if dim_second_last == hidden_size: + # w13 is [E, hidden_size, 2*inter_size] - standard layout + inter_size = dim_last // 2 + needs_transpose = False + else: + # w13 is [E, 2*inter_size, hidden_size] - needs transpose + inter_size = dim_second_last // 2 + needs_transpose = True + + assert w13.is_contiguous() and w2.is_contiguous() + + # 4bits support [E, N, K] + # other types [E, K, N] + if not is_int4 and not is_mxfp4 and not is_mxfp8: + if not hasattr(w13, 'xpu_fused_moe'): + if needs_transpose: + w13.data = w13.transpose(-1, -2).contiguous() + w2.data = w2.transpose(-1, -2).contiguous() + w13.xpu_fused_moe = True + w13.inter_size = inter_size + else: + inter_size = w13.inter_size + + if is_int4 and not hasattr(w13, 'xpu_fused_moe'): + w13_tmp = torch.empty_like(w13) + w2_tmp = torch.empty_like(w2) + for i in range(num_experts): + w13_tmp[i] = implement_zp(w13[i]) + w2_tmp[i] = implement_zp(w2[i]) + w13_tmp = w13_tmp.contiguous() + w2_tmp = w2_tmp.contiguous() + w13.data = w13_tmp + w2.data = w2_tmp + w13.xpu_fused_moe = True + + # TODO: will all integrated in Cpp func. Temporary expose before gemm fusion + num_experts_per_node = num_experts + experts_per_token = n_experts_per_token + num_moe_inputs = n_experts_per_token * num_rows + permuted_elems = num_moe_inputs * hidden_size + # interbuf_elems = num_moe_inputs * inter_size + permuted_row_to_unpermuted_row_size = num_moe_inputs * 4 + permuted_token_selected_experts_size = num_moe_inputs * 4 + src_to_dest_map_size = experts_per_token * num_rows * 4 + expert_first_token_offset_size = (num_experts_per_node + 1) * 8 + num_tokens_per_block = compute_num_tokens_per_block( + num_rows, num_experts_per_node) + num_blocks_per_seq = ceilDiv(num_rows, num_tokens_per_block) + blocked_expert_counts_size = num_experts_per_node * num_blocks_per_seq * 4 + blocked_expert_counts_cumsum_size = blocked_expert_counts_size + blocked_row_to_unpermuted_row_size = num_experts_per_node * num_rows * 4 + permuted_data_size = permuted_elems * hidden_states.element_size() + permuted_token_final_scales_size = num_moe_inputs * 4 + + ws_map = {} + map_offset = 0 + + def config_ws(name, size): + nonlocal map_offset + if size % 256 != 0: + size += 256 - size % 256 + ws_map[name] = (size, map_offset) + map_offset += size + + config_ws("permuted_row_to_unpermuted_row", + permuted_row_to_unpermuted_row_size) + config_ws("permuted_token_selected_experts", + permuted_token_selected_experts_size) + config_ws("unpermuted_row_to_permuted_row", src_to_dest_map_size) + config_ws("blocked_expert_counts", blocked_expert_counts_size) + config_ws("blocked_expert_counts_cumsum", + blocked_expert_counts_cumsum_size) + config_ws("blocked_row_to_unpermuted_row", + blocked_row_to_unpermuted_row_size) + config_ws("expert_first_token_offset", expert_first_token_offset_size) + config_ws("permuted_token_final_scales", permuted_token_final_scales_size) + config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size) + + workspace = torch.zeros(map_offset, + dtype=torch.uint8, + device=hidden_states.device) + if topk_ids.dtype == torch.int32: + topk_ids = topk_ids.to(torch.int64) + ops.fused_moe_prologue( + input=hidden_states, + token_selected_experts=topk_ids, + token_final_scales=topk_weights, + workspace=workspace, + hidden_size=hidden_size, + inter_size=inter_size, + ep_rank=ep_rank, + ep_size=ep_size, + num_experts_on_rank=num_experts_per_node) + + expert_first_token_offset_bytes = workspace[ + ws_map["expert_first_token_offset"][1]: + ws_map["expert_first_token_offset"][1] + + expert_first_token_offset_size] + unpermuted_row_to_permuted_row_bytes = workspace[ + ws_map["unpermuted_row_to_permuted_row"][1]: + ws_map["unpermuted_row_to_permuted_row"][1] + + src_to_dest_map_size] + permuted_row_to_unpermuted_row_bytes = workspace[ + ws_map["permuted_row_to_unpermuted_row"][1]: + ws_map["permuted_row_to_unpermuted_row"][1] + + permuted_row_to_unpermuted_row_size] + + if torch.compiler.is_compiling(): + expert_first_token_offset = _bytes_to_typed_tensor( + expert_first_token_offset_bytes, torch.int64 + ) + unpermuted_row_to_permuted_row = _bytes_to_typed_tensor( + unpermuted_row_to_permuted_row_bytes, torch.int32 + ) + permuted_row_to_unpermuted_row = _bytes_to_typed_tensor( + permuted_row_to_unpermuted_row_bytes, torch.int32 + ) + else: + expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64) + unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32) + permuted_row_to_unpermuted_row = permuted_row_to_unpermuted_row_bytes.view(torch.int32) + gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]: + ws_map["overlapped_gemm1_gemm2_inputs"][1] + + permuted_data_size].view(hidden_states.dtype).view( + num_moe_inputs, hidden_size) + # permuted_token_final_scales = workspace[ + # ws_map["permuted_token_final_scales"][1]: + # ws_map["permuted_token_final_scales"][1] + + # permuted_token_final_scales_size].view(torch.float) + gemm1_output = torch.empty((num_moe_inputs, 2 * inter_size), + dtype=hidden_states.dtype, + device=hidden_states.device) + + ########### gemm1 ################## + input_B = w13 + + if not is_fp8 and not is_int4 and not is_mxfp4 and not is_mxfp8: + ops.cutlass_grouped_gemm_interface( + ptr_A=gemm1_input, + ptr_B=input_B, + ptr_scales=None, + ptr_bias=w13_bias, + ptr_D=gemm1_output, + expert_first_token_offset=expert_first_token_offset, + N=2 * inter_size, + K=hidden_size, + num_experts=num_experts_per_node, + is_B_int4=is_int4, + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) + else: + ops.cutlass_grouped_gemm_interface( + ptr_A=gemm1_input, + ptr_B=input_B, + ptr_scales=w13_scales, + ptr_bias=w13_bias, + ptr_D=gemm1_output, + expert_first_token_offset=expert_first_token_offset, + N=2 * inter_size, + K=hidden_size, + num_experts=num_experts_per_node, + is_B_int4=is_int4, + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) + + # act + act_output = torch.empty((num_moe_inputs, inter_size), + dtype=gemm1_output.dtype, + device=gemm1_output.device) + if activation == "silu": + ops.silu_and_mul(act_output, gemm1_output) + elif activation == "gelu": + ops.gelu_and_mul(act_output, gemm1_output) + elif activation == "swigluoai": + ops.swigluoai_and_mul(act_output, gemm1_output, 1.702, 7.0) + else: + raise ValueError(f"Unsupported FusedMoe activation: {activation}.") + + ########### gemm2 ################## + input_A = act_output.contiguous() + input_B = w2 + gemm2_output = torch.empty((num_moe_inputs, hidden_size), + dtype=hidden_states.dtype, + device=hidden_states.device) + if not is_fp8 and not is_int4 and not is_mxfp4 and not is_mxfp8: + ops.cutlass_grouped_gemm_interface( + ptr_A=input_A, + ptr_B=input_B, + ptr_scales=None, + ptr_bias=w2_bias, + ptr_D=gemm2_output, + expert_first_token_offset=expert_first_token_offset, + N=hidden_size, + K=inter_size, + num_experts=num_experts_per_node, + is_B_int4=is_int4, + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) + else: + ops.cutlass_grouped_gemm_interface( + ptr_A=input_A, + ptr_B=input_B, + ptr_scales=w2_scales, + ptr_bias=w2_bias, + ptr_D=gemm2_output, + expert_first_token_offset=expert_first_token_offset, + N=hidden_size, + K=inter_size, + num_experts=num_experts_per_node, + is_B_int4=is_int4, + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) + + ops.moe_gather(output, gemm2_output, topk_weights, + permuted_row_to_unpermuted_row, + unpermuted_row_to_permuted_row, + expert_first_token_offset, + num_experts_per_node) + return output + + +def apply_jitter(x: torch.Tensor, moe_jitter_eps: float) -> torch.Tensor: + """Apply jitter to the input tensor for regularization.""" + low = 1.0 - moe_jitter_eps + high = 1.0 + moe_jitter_eps + noise = torch.rand(x.size(), dtype=x.dtype, device=x.device) + return x * (low + noise * (high - low)) + + +def compute_top_k(scores: torch.Tensor, moe_top_k: int): + """Compute the top-k scores from the logits.""" + if moe_top_k == 1: + return scores.max(dim=-1, keepdim=True) + return torch.topk(scores, moe_top_k, dim=-1) + + +def route_tokens_xpu( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + moe_top_k: int, + moe_num_experts: int, + moe_jitter_eps: float = None, + moe_normalize_expert_weights: int = None, + training: bool = False, +) -> tuple: + """Route tokens to experts and compute expert weights and indices (XPU version).""" + if training and moe_jitter_eps is not None: + x = apply_jitter(x, moe_jitter_eps) + + x_flat = x.view(-1, x.shape[-1]) + logits = torch.nn.functional.linear(x_flat, router_weight, router_bias) + expert_weights, expert_indices = compute_top_k(logits, moe_top_k) + expert_weights = expert_weights.softmax(dim=-1) + if moe_normalize_expert_weights is not None: + expert_weights = expert_weights / torch.norm( + expert_weights, + p=moe_normalize_expert_weights, + dim=-1, + keepdim=True, + ) + + return logits, expert_weights, expert_indices + + +def _get_device_mesh(model): + """Extract device_mesh from child's unused pre_hook closure for EP support.""" + try: + hook = next( + h + for h in model.experts._forward_pre_hooks.values() + if "device_mesh" in h.__code__.co_freevars + ) + return hook.__closure__[ + hook.__code__.co_freevars.index("device_mesh") + ].cell_contents + except Exception: + return None + + +def _mxfp4_expert_tensors(experts, dtype): + """View transformers' MXFP4 experts as the tensors the SYCL grouped GEMM expects. + + Weights and scales arrive wrapped in triton tensors. On XPU both use StridedLayout, + whose swizzle is the identity, so `storage.data` is still the checkpoint tensor - just + transposed to [E, K/2, N]. Transposing back therefore restores the original contiguous + [E, N, K/2] weights and [E, N, K/32] E8M0 scales as plain views: `contiguous()` finds + them already contiguous and returns them untouched, so this costs no copy and no extra + memory. The experts module is left alone, which keeps the triton tensors and the + PrecisionConfig that `save_pretrained` needs to re-serialize the checkpoint. + """ + tensors = [] + for proj in ("gate_up_proj", "down_proj"): + config = getattr(experts, f"{proj}_precision_config") + tensors.append(getattr(experts, proj).storage.data.transpose(-1, -2).contiguous()) + tensors.append(config.weight_scale.storage.data.transpose(-1, -2).contiguous()) + + # The GEMM reads the bias as ElementA; `to` is a no-op when it already matches. + bias = getattr(experts, f"{proj}_bias", None) + tensors.append(None if bias is None else bias.data.to(dtype)) + + return tensors + + +class MegaBlocksMoeMLP(torch.nn.Module): + can_torch_compile: bool = True + + def forward(self, x: torch.Tensor) -> tuple: + """ + Forward pass through the MoE layer. + + Args: + x: Input tensor of shape [batch_size, seq_len, hidden_size] or [tokens, hidden_size] + + Returns: + Tuple of (output, expert_weights) where: + - output: Tensor of same shape as input + - expert_weights: Expert weights for each token [tokens, top_k] + """ + + # Get MoE parameters from the wrapped modules + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + moe_jitter_eps = getattr(self.experts, "jitter_eps", None) + moe_normalize_expert_weights = getattr( + self.experts, "normalize_expert_weights", None + ) + + # Get EP (Expert Parallelism) parameters + ep_size = 1 + ep_rank = 0 + expert_parallel_group = getattr(self, "expert_parallel_group", None) + if expert_parallel_group is None: + device_mesh = _get_device_mesh(self) + if device_mesh is not None: + expert_parallel_group = device_mesh.get_group() + if expert_parallel_group is not None: + import torch.distributed as dist + if dist.is_initialized(): + ep_size = dist.get_world_size(expert_parallel_group) + ep_rank = dist.get_rank(expert_parallel_group) + + # Number of experts on this rank + num_experts_on_rank = moe_num_experts // ep_size + + # Detect activation type - check for GptOss-style swigluoai activation + # GptOssExperts has alpha and limit attributes for swigluoai + if hasattr(self.experts, "alpha") and hasattr(self.experts, "limit"): + activation = "swigluoai" + else: + activation = getattr(self.experts, "activation", "silu") + + # Get weight tensors - support different naming conventions + if hasattr(self.experts, "gate_up_proj"): + w13 = self.experts.gate_up_proj + # NOTE: swigluoai_and_mul kernel expects interleaved layout [g0,u0,g1,u1,...] + # which matches GptOss's gate_up_proj layout, so no conversion needed. + + elif hasattr(self.experts, "w1"): + # Combine w1 and w3 if stored separately + w1 = self.experts.w1 + w3 = getattr(self.experts, "w3", None) + if w3 is not None: + w13 = torch.cat([w1, w3], dim=-2) + else: + w13 = w1 + else: + raise AttributeError("experts module must have 'gate_up_proj' or 'w1' attribute") + + if hasattr(self.experts, "down_proj"): + w2 = self.experts.down_proj + elif hasattr(self.experts, "w2"): + w2 = self.experts.w2 + else: + raise AttributeError("experts module must have 'down_proj' or 'w2' attribute") + + # Get optional bias tensors + w13_bias = getattr(self.experts, "gate_up_proj_bias", None) + w2_bias = getattr(self.experts, "down_proj_bias", None) + + # Get quantization info + is_fp8 = getattr(self.experts, "is_fp8", False) + is_int4 = getattr(self.experts, "is_int4", False) + is_mxfp4 = getattr(self.experts, "is_mxfp4", False) + is_mxfp8 = getattr(self.experts, "is_mxfp8", False) + + w13_scales = getattr(self.experts, "gate_up_proj_scales", None) + w2_scales = getattr(self.experts, "down_proj_scales", None) + + # Native MXFP4 checkpoint: derive the GEMM's view of the experts once and cache it. + mxfp4 = getattr(self, "_mxfp4_tensors", None) + if (mxfp4 is None or mxfp4[0] is not x.dtype) and getattr( + self.experts, "gate_up_proj_precision_config", None + ) is not None: + mxfp4 = (x.dtype, *_mxfp4_expert_tensors(self.experts, x.dtype)) + self._mxfp4_tensors = mxfp4 + if mxfp4 is not None: + _, w13, w13_scales, w13_bias, w2, w2_scales, w2_bias = mxfp4 + is_mxfp4 = True + + # Store original shape + in_shape = x.size() + + # Route tokens to experts + logits, expert_weights, expert_indices = route_tokens_xpu( + x, + self.router.weight, + getattr(self.router, "bias", None), + moe_top_k, + moe_num_experts, + moe_jitter_eps, + moe_normalize_expert_weights, + self.training, + ) + + # Reshape input for fused MoE + x_flat = x.view(-1, x.shape[-1]) + + # Call XPU fused MoE kernel + output = xpu_fused_moe( + hidden_states=x_flat, + w13=w13, + w13_scales=w13_scales, + w13_bias=w13_bias, + w2=w2, + w2_scales=w2_scales, + w2_bias=w2_bias, + topk_weights=expert_weights.float(), + topk_ids=expert_indices, + n_experts_per_token=moe_top_k, + activation=activation, + num_experts=num_experts_on_rank, + ep_rank=ep_rank, + ep_size=ep_size, + is_fp8=is_fp8, + is_int4=is_int4, + is_mxfp4=is_mxfp4, + is_mxfp8=is_mxfp8, + ) + + # All-reduce across EP group to combine partial expert outputs + if ep_size > 1 and expert_parallel_group is not None: + import torch.distributed as dist + dist.all_reduce(output, op=dist.ReduceOp.SUM, group=expert_parallel_group) + + # Restore original shape + output = output.view(in_shape) + + return output, expert_weights + + +# Export classes and functions +__all__ = [ + "MegaBlocksMoeMLP", + "xpu_fused_moe", + "cutlass_grouped_gemm", + "cutlass_grouped_gemm_xe2", +] \ No newline at end of file diff --git a/build/torch214-cxx11-cu130-x86_64-linux/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0766d7b8da4f97baca212177b4bb989bc6374bf8 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/__init__.py @@ -0,0 +1,205 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch + +# Stable alias: bare `ops` is shadowed by `from . import layers` below. +from ._ops import ops as _compiled_ops +from . import ops + +from .grouped_gemm import backend as gg_backend +from .grouped_gemm import ops as gg_ops + + +from ._layers.arguments import Arguments +from ._layers.dmoe import ParallelDroplessMLP, dMoE +from ._layers.glu import SparseGLU +from ._layers.mlp import MLP, SparseMLP +from ._layers.moe import MoE, ParallelMLP, get_load_balancing_loss + +from . import layers + +# This section contains the direct kernel exports (not inlcuded in the original code) +def exclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor: + """ + Compute exclusive cumulative sum along the specified dimension. + + Args: + x: Input tensor + dim: Dimension along which to compute cumsum + out: Output tensor (modified in-place) + + Returns: + The output tensor + """ + result = ops.exclusive_cumsum(x, dim) + out.copy_(result) + return out + + +def inclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor: + """ + Compute inclusive cumulative sum along the specified dimension. + + Args: + x: Input tensor + dim: Dimension along which to compute cumsum + out: Output tensor (modified in-place) + + Returns: + The output tensor + """ + result = ops.inclusive_cumsum(x, dim) + out.copy_(result) + return out + + +def histogram(x: torch.Tensor, num_bins: int) -> torch.Tensor: + """ + Compute histogram of input tensor values. + + Args: + x: Input tensor + num_bins: Number of histogram bins + + Returns: + Histogram tensor with counts for each bin + """ + return ops.histogram(x, num_bins) + + +def indices( + padded_bins: torch.Tensor, + block_size: int, + output_block_rows: int, + output_block_columns: int, +) -> torch.Tensor: + """ + Construct indices from padded bins for sparse operations. + + Args: + padded_bins: Tensor containing bin boundaries + block_size: Size of each block + output_block_rows: Number of rows in output blocks + output_block_columns: Number of columns in output blocks + + Returns: + Tensor containing constructed indices + """ + return ops.indices(padded_bins, block_size, output_block_rows, output_block_columns) + + +def replicate_forward( + x: torch.Tensor, bins: torch.Tensor, out: torch.Tensor +) -> torch.Tensor: + """ + Forward pass of replicate operation - replicate values according to bin sizes. + + Args: + x: Input tensor with values to replicate + bins: Tensor containing bin sizes + out: Output tensor (modified in-place) + + Returns: + The output tensor + """ + return ops.replicate_forward(x, bins, out) + + +def replicate_backward( + grad: torch.Tensor, bins: torch.Tensor, out: torch.Tensor +) -> torch.Tensor: + """ + Backward pass of replicate operation - reduce gradients back to bins. + + Args: + grad: Gradient tensor to reduce + bins: Tensor containing bin sizes + out: Output tensor (modified in-place) + + Returns: + The output tensor + """ + return ops.replicate_backward(grad, bins, out) + + +def sort( + x: torch.Tensor, end_bit: int, x_out: torch.Tensor, iota_out: torch.Tensor +) -> torch.Tensor: + """ + Radix sort with index tracking. + + Args: + x: Input tensor to sort + end_bit: Number of bits to consider in sorting + x_out: Output tensor for sorted values + iota_out: Output tensor for sorted indices + + Returns: + The sorted values tensor + """ + _compiled_ops.sort(x, end_bit, x_out, iota_out) + return x_out + + +# Convenience functions for common use cases +def cumsum(x: torch.Tensor, dim: int = -1, exclusive: bool = False) -> torch.Tensor: + """ + Compute cumulative sum with automatic output allocation. + + Args: + x: Input tensor + dim: Dimension along which to compute cumsum (default: last dimension) + exclusive: Whether to compute exclusive (True) or inclusive (False) cumsum + + Returns: + New tensor containing the cumulative sum + """ + out = torch.empty_like(x) + if exclusive: + return exclusive_cumsum(x, dim, out) + else: + return inclusive_cumsum(x, dim, out) + + +def argsort(x: torch.Tensor, end_bit: int = 32) -> tuple[torch.Tensor, torch.Tensor]: + """ + Sort tensor and return both sorted values and indices. + + Args: + x: Input tensor to sort + end_bit: Number of bits to consider in sorting + + Returns: + Tuple of (sorted_values, sorted_indices) + """ + x_out = torch.empty_like(x) + iota_out = torch.empty_like(x) + sort(x, end_bit, x_out, iota_out) + return x_out, iota_out + + +# Export public API +__all__ = [ + "MyReplacementLayer", + # Direct kernel exports + "exclusive_cumsum", + "inclusive_cumsum", + "histogram", + "indices", + "replicate_forward", + "replicate_backward", + "sort", + "cumsum", + "argsort", + # Original exports + "Arguments", + "ParallelDroplessMLP", + "dMoE", + "SparseGLU", + "MLP", + "SparseMLP", + "MoE", + "ParallelMLP", + "get_load_balancing_loss", +] diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/_triton_kernels/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/_triton_kernels/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/_triton_kernels/gmm.py b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/_triton_kernels/gmm.py new file mode 100644 index 0000000000000000000000000000000000000000..8c101d07cea416f9390b708e5a35fdc466e48aed --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/_triton_kernels/gmm.py @@ -0,0 +1,574 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. + + +# Imports. +# ------------------------------------------------------------------------------ + +# Python standard library +import functools + +# Triton +import triton +import triton.language as tl + +# AITER +from ..configs import CONFIGS as _CONFIGS +from ..utils._triton import arch_info +from ..utils._triton.pid_preprocessing import pid_grid, remap_xcd + +# Kernel config. +# ------------------------------------------------------------------------------ + + +@functools.lru_cache() +def get_config( + gmm_type: str, M: int, K: int, N: int, G: int, accumulate: bool = False +) -> dict[str, int]: + assert gmm_type in { + "gmm", + "ptgmm", + "nptgmm", + }, f"'{gmm_type}' is an invalid GMM variant." + dev = arch_info.get_arch() + assert ( + dev in _CONFIGS + ), f"No GMM configuration tuned for arch '{dev}'. Supported: {sorted(_CONFIGS)}." + arch_configs = _CONFIGS[dev] + assert ( + "default" in arch_configs[gmm_type] + ), "Default configuration is absent." + key = "accumulate" if accumulate else "default" + return arch_configs[gmm_type][key] + + +# Common code shared by GMM and TGMM kernels. +# ------------------------------------------------------------------------------ + + +# XCD remapping followed by 1D PID to 2D grid mapping. +@triton.jit +def _remap_xcd_tile_grid( + tile_in_mm, + num_row_tiles, + num_col_tiles, + GROUP_SIZE: tl.constexpr = 1, + NUM_XCDS: tl.constexpr = 8, +): + return pid_grid( + remap_xcd(tile_in_mm, num_row_tiles * num_col_tiles, NUM_XCDS=NUM_XCDS), + num_row_tiles, + num_col_tiles, + GROUP_SIZE_M=GROUP_SIZE, + ) + + +# GMM kernel. +# ------------------------------------------------------------------------------ + + +@triton.heuristics( + { + "K_DIVISIBLE_BY_BLOCK_SIZE_K": lambda META: META["K"] % META["BLOCK_SIZE_K"] + == 0, + } +) +@triton.jit +def gmm_kernel( + # Tensor pointers: + lhs_ptr, + rhs_ptr, + group_sizes_ptr, + out_ptr, + bias_ptr, + # Tensor shapes: + M: int, + K: int, + N: int, + G: int, + # Meta-parameters: + TRANS_RHS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + K_DIVISIBLE_BY_BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE: tl.constexpr, + GRID_DIM: tl.constexpr, + USE_BIAS: tl.constexpr, +): + tl.assume(M > 0) + tl.assume(K > 0) + tl.assume(N > 0) + tl.assume(G > 0) + + num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) + tl.device_assert(num_n_tiles > 0, "num_n_tiles <= 0") + + # Current tile. Each program computes multiple tiles of each group. + tile = tl.program_id(0) + tl.device_assert(tile >= 0, "tile < 0 (at initialization)") + + # Tile limit of last MM problem (inclusive). + last_mm_tile = 0 + + # Last input row of lhs and output row of out. Each group reads some rows of + # lhs and writes some rows to out. + last_m = 0 + + # Loop through all (m, K, N) MM problems: + # (m, K) x (K, N) = (m, N) + # sum(m) = M + for g in range(G): + # Get m dimension of current MM problem. + m = tl.load(group_sizes_ptr + g) + # m can be zero if group is empty + tl.device_assert(m >= 0, "m < 0") + + num_m_tiles = tl.cdiv(m, BLOCK_SIZE_M) + # num_m_tiles can be zero if group is empty + tl.device_assert(num_m_tiles >= 0, "num_m_tiles < 0") + + num_tiles = num_m_tiles * num_n_tiles + # num_tiles can be zero if group is empty + tl.device_assert(num_tiles >= 0, "num_tiles < 0") + + # Loop through tiles of current MM problem. + while tile >= last_mm_tile and tile < last_mm_tile + num_tiles: + # Figure out tile coordinates in current MM problem. + tile_in_mm = tile - last_mm_tile + tl.device_assert(tile_in_mm >= 0, "tile_in_mm < 0") + + tile_m, tile_n = _remap_xcd_tile_grid( + tile_in_mm, num_m_tiles, num_n_tiles, GROUP_SIZE=GROUP_SIZE + ) + + # Do regular MM: + + tl.device_assert(tile_m * BLOCK_SIZE_M >= 0, "tile_m * BLOCK_SIZE_M < 0") + tl.device_assert(tile_n * BLOCK_SIZE_N >= 0, "tile_n * BLOCK_SIZE_N < 0") + + offs_lhs_m = ( + tile_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + ) % m + offs_rhs_n = ( + tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + ) % N + offs_k = tl.arange(0, BLOCK_SIZE_K).to(tl.int64) + + lhs_ptrs = lhs_ptr + (last_m + offs_lhs_m[:, None]) * K + offs_k[None, :] + + if TRANS_RHS: + rhs_ptrs = ( + rhs_ptr + + g.to(tl.int64) * K * N + + offs_k[:, None] + + offs_rhs_n[None, :] * K + ) + else: + rhs_ptrs = ( + rhs_ptr + + g.to(tl.int64) * K * N + + offs_k[:, None] * N + + offs_rhs_n[None, :] + ) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + if K_DIVISIBLE_BY_BLOCK_SIZE_K: + lhs = tl.load(lhs_ptrs) + rhs = tl.load(rhs_ptrs) + else: + k_mask_limit = K - k * BLOCK_SIZE_K + lhs = tl.load( + lhs_ptrs, mask=offs_k[None, :] < k_mask_limit, other=0 + ) + rhs = tl.load( + rhs_ptrs, mask=offs_k[:, None] < k_mask_limit, other=0 + ) + + acc = tl.dot(lhs, rhs, acc=acc) + + lhs_ptrs += BLOCK_SIZE_K + + if TRANS_RHS: + rhs_ptrs += BLOCK_SIZE_K + else: + rhs_ptrs += BLOCK_SIZE_K * N + + # Add bias if enabled + if USE_BIAS: + offs_bias_n = tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange( + 0, BLOCK_SIZE_N + ) + bias_ptrs = bias_ptr + g.to(tl.int64) * N + offs_bias_n + bias = tl.load(bias_ptrs, mask=offs_bias_n < N, other=0.0) + # Convert bias to float32 to match accumulator precision + bias = bias.to(tl.float32) + # Broadcast bias across M dimension and add in float32 + acc += bias[None, :] + + # Convert to output dtype after all computations + acc = acc.to(out_ptr.type.element_ty) + + offs_out_m = tile_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_out_n = tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + + out_ptrs = ( + out_ptr + (last_m + offs_out_m[:, None]) * N + offs_out_n[None, :] + ) + + tl.store( + out_ptrs, + acc, + mask=(offs_out_m[:, None] < m) & (offs_out_n[None, :] < N), + ) + + # Go to the next tile by advancing number of programs. + tile += GRID_DIM + tl.device_assert(tile > 0, "tile <= 0 (at update)") + + # Get ready to go to the next MM problem. + + last_mm_tile += num_tiles + # last_mm_tile can be zero if group 0 is skipped + tl.device_assert(last_mm_tile >= 0, "last_mm_tile < 0 (at update)") + + last_m += m + # last_m can be zero if group 0 is skipped + tl.device_assert(last_m >= 0, "last_m < 0 (at update)") + tl.device_assert(last_m <= M, "last_m > M (at update)") + + +# Persistent TGMM kernel. +# ------------------------------------------------------------------------------ + + +@triton.jit +def tgmm_persistent_kernel( + # Tensor pointers: + lhs_ptr, + rhs_ptr, + group_sizes_ptr, + out_ptr, + bias_grad_ptr, + # Tensor shapes: + M: int, + K: int, + N: int, + G: int, + # Meta-parameters: + TRANS_LHS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + GROUP_SIZE: tl.constexpr, + GRID_DIM: tl.constexpr, + COMPUTE_BIAS_GRAD: tl.constexpr, + ACCUMULATE: tl.constexpr, +): + tl.assume(M > 0) + tl.assume(K > 0) + tl.assume(N > 0) + tl.assume(G > 0) + + num_k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + tl.device_assert(num_k_tiles > 0, "num_k_tiles <= 0") + + num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) + tl.device_assert(num_n_tiles > 0, "num_n_tiles <= 0") + + num_tiles = num_k_tiles * num_n_tiles + tl.device_assert(num_tiles > 0, "num_tiles <= 0") + + # Current tile. Each program computes multiple tiles of each group. + tile = tl.program_id(0) + tl.device_assert(tile >= 0, "tile < 0 (at initialization)") + + # Tile limit of last MM problem (inclusive). + last_mm_tile = 0 + + # Last input column of lhs and input row of rhs. Each group reads some + # columns of lhs and some rows of rhs. + last_m = 0 + + # Loop through all (K, m, N) MM problems: + # (K, m) x (m, N) = (K, N) + # sum(m) = M + for g in range(G): + # Get m dimension of current MM problem. + m = tl.load(group_sizes_ptr + g) + # m can be zero if group is empty + tl.device_assert(m >= 0, "m < 0") + + # Loop through tiles of current MM problem. + while tile >= last_mm_tile and tile < last_mm_tile + num_tiles: + # Figure out tile coordinates in current MM problem. + tile_in_mm = tile - last_mm_tile + tl.device_assert(tile_in_mm >= 0, "tile_in_mm < 0") + + tile_k, tile_n = _remap_xcd_tile_grid( + tile_in_mm, num_k_tiles, num_n_tiles, GROUP_SIZE=GROUP_SIZE + ) + + # Do regular MM: + + tl.device_assert(tile_k * BLOCK_SIZE_K >= 0, "tile_k * BLOCK_SIZE_K < 0") + tl.device_assert(tile_n * BLOCK_SIZE_N >= 0, "tile_n * BLOCK_SIZE_N < 0") + + offs_lhs_k = ( + tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + ) % K + offs_rhs_n = ( + tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + ) % N + offs_m = tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + + if TRANS_LHS: + lhs_ptrs = ( + lhs_ptr + offs_lhs_k[:, None] + (last_m + offs_m[None, :]) * K + ) + else: + lhs_ptrs = ( + lhs_ptr + offs_lhs_k[:, None] * M + (last_m + offs_m[None, :]) + ) + + rhs_ptrs = rhs_ptr + (last_m + offs_m[:, None]) * N + offs_rhs_n[None, :] + + loop_m = tl.cdiv(m, BLOCK_SIZE_M) + m_divisible_by_block_m = m % BLOCK_SIZE_M == 0 + if not m_divisible_by_block_m: + loop_m -= 1 + + acc = tl.zeros((BLOCK_SIZE_K, BLOCK_SIZE_N), dtype=tl.float32) + + # Initialize bias accumulator + bias_acc = tl.zeros((BLOCK_SIZE_K,), dtype=tl.float32) + + for _ in range(0, loop_m): + lhs = tl.load(lhs_ptrs) + rhs = tl.load(rhs_ptrs) + + acc = tl.dot(lhs, rhs, acc=acc) + + # Accumulate for bias gradient: sum lhs across M dimension + if COMPUTE_BIAS_GRAD and tile_n == 0: + bias_acc += tl.sum( + lhs, axis=1 + ) # Sum across M dimension [K, M] -> [K] + + if TRANS_LHS: + lhs_ptrs += BLOCK_SIZE_M * K + else: + lhs_ptrs += BLOCK_SIZE_M + + rhs_ptrs += BLOCK_SIZE_M * N + + if not m_divisible_by_block_m: + offs_lhs_k = ( + tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + ) % K + offs_rhs_n = ( + tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + ) % N + offs_m = loop_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + lhs = tl.load(lhs_ptrs, mask=offs_m[None, :] < m, other=0) + rhs = tl.load(rhs_ptrs, mask=offs_m[:, None] < m, other=0) + acc = tl.dot(lhs, rhs, acc=acc) + + # Accumulate last chunk for bias gradient + if COMPUTE_BIAS_GRAD and tile_n == 0: + bias_acc += tl.sum(lhs, axis=1) + + acc = acc.to(out_ptr.type.element_ty) + + offs_out_k = tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + offs_out_n = tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + + out_ptrs = ( + out_ptr + + g.to(tl.int64) * K * N + + offs_out_k[:, None] * N + + offs_out_n[None, :] + ) + + mask = (offs_out_k[:, None] < K) & (offs_out_n[None, :] < N) + if ACCUMULATE: + # Load existing values and add to them (like beta=1 in BLAS) + old_vals = tl.load(out_ptrs, mask=mask, other=0.0) + tl.store(out_ptrs, acc + old_vals, mask=mask) + else: + # Overwrite output (like beta=0 in BLAS) + tl.store(out_ptrs, acc, mask=mask) + + # Store bias gradient (only for first N tile, sum across all M) + if COMPUTE_BIAS_GRAD and tile_n == 0: + # Keep as float32 for atomic_add (bf16 not supported for atomics) + bias_grad_ptrs = bias_grad_ptr + g.to(tl.int64) * K + offs_out_k + # Use atomic add since multiple K-tiles may write to same expert's bias + tl.atomic_add( + bias_grad_ptrs, bias_acc, mask=offs_out_k < K, sem="relaxed" + ) + + # Go to the next tile by advancing number of programs. + tile += GRID_DIM + tl.device_assert(tile > 0, "tile <= 0 (at update)") + + # Get ready to go to the next MM problem. + + last_mm_tile += num_tiles + # last_mm_tile can be zero if group 0 is skipped + tl.device_assert(last_mm_tile >= 0, "last_mm_tile < 0 (at update)") + + last_m += m + # last_m can be zero if group 0 is skipped + tl.device_assert(last_m >= 0, "last_m < 0 (at update)") + tl.device_assert(last_m <= M, "last_m > M (at update)") + + +# Regular non-persistent TGMM kernel. +# ------------------------------------------------------------------------------ + + +@triton.heuristics({"BLOCK_SIZE_G": lambda META: triton.next_power_of_2(META["G"])}) +@triton.jit +def tgmm_non_persistent_kernel( + # Tensor pointers: + lhs_ptr, + rhs_ptr, + group_sizes_ptr, + out_ptr, + bias_grad_ptr, + # Tensor shapes: + M: int, + K: int, + N: int, + G: int, + # Meta-parameters: + TRANS_LHS: tl.constexpr, + BLOCK_SIZE_G: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + GROUP_SIZE: tl.constexpr, + COMPUTE_BIAS_GRAD: tl.constexpr, + ACCUMULATE: tl.constexpr, +): + tl.assume(M > 0) + tl.assume(K > 0) + tl.assume(N > 0) + tl.assume(G > 0) + + # Get group ID from grid. + g = tl.program_id(0) + tl.device_assert(g >= 0, "g < 0") + tl.device_assert(g < G, "g >= G") + + # Get m dimension of current MM group. + m = tl.load(group_sizes_ptr + g) + # m can be zero if group is empty. + tl.device_assert(m >= 0, "m < 0") + + # Skip empty groups. + if m == 0: + return + + # Compute sum(group_sizes) until current group g. + # It's the starting column of lhs and starting row of rhs. + offs_g = tl.arange(0, BLOCK_SIZE_G) + group_sizes = tl.load(group_sizes_ptr + offs_g, mask=offs_g < g, other=0) + start_m = tl.sum(group_sizes) + + num_k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + tl.device_assert(num_k_tiles > 0, "num_k_tiles <= 0") + + num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) + tl.device_assert(num_n_tiles > 0, "num_n_tiles <= 0") + + # Get MM tile from grid. + tile_in_mm = tl.program_id(1) + tl.device_assert(tile_in_mm >= 0, "tile_in_mm < 0") + + tile_k, tile_n = _remap_xcd_tile_grid( + tile_in_mm, num_k_tiles, num_n_tiles, GROUP_SIZE=GROUP_SIZE + ) + + tl.device_assert(tile_k * BLOCK_SIZE_K >= 0, "tile_k * BLOCK_SIZE_K < 0") + tl.device_assert(tile_n * BLOCK_SIZE_N >= 0, "tile_n * BLOCK_SIZE_N < 0") + + offs_lhs_k = (tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)) % K + offs_rhs_n = (tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_m = tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + + if TRANS_LHS: + lhs_ptrs = lhs_ptr + offs_lhs_k[:, None] + (start_m + offs_m[None, :]) * K + else: + lhs_ptrs = lhs_ptr + offs_lhs_k[:, None] * M + (start_m + offs_m[None, :]) + + rhs_ptrs = rhs_ptr + (start_m + offs_m[:, None]) * N + offs_rhs_n[None, :] + + loop_m = tl.cdiv(m, BLOCK_SIZE_M) + m_divisible_by_block_m = m % BLOCK_SIZE_M == 0 + if not m_divisible_by_block_m: + loop_m -= 1 + + acc = tl.zeros((BLOCK_SIZE_K, BLOCK_SIZE_N), dtype=tl.float32) + # Initialize bias accumulator + bias_acc = tl.zeros((BLOCK_SIZE_K,), dtype=tl.float32) + + for _ in range(0, loop_m): + lhs = tl.load(lhs_ptrs) + rhs = tl.load(rhs_ptrs) + + acc = tl.dot(lhs, rhs, acc=acc) + + # Accumulate for bias gradient: sum lhs across M dimension + if COMPUTE_BIAS_GRAD and tile_n == 0: + bias_acc += tl.sum(lhs, axis=1) # [K, M] -> [K] + + if TRANS_LHS: + lhs_ptrs += BLOCK_SIZE_M * K + else: + lhs_ptrs += BLOCK_SIZE_M + + rhs_ptrs += BLOCK_SIZE_M * N + + if not m_divisible_by_block_m: + offs_lhs_k = ( + tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + ) % K + offs_rhs_n = ( + tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + ) % N + offs_m = loop_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + lhs = tl.load(lhs_ptrs, mask=offs_m[None, :] < m, other=0) + rhs = tl.load(rhs_ptrs, mask=offs_m[:, None] < m, other=0) + acc = tl.dot(lhs, rhs, acc=acc) + # Accumulate last chunk for bias gradient + if COMPUTE_BIAS_GRAD and tile_n == 0: + bias_acc += tl.sum(lhs, axis=1) + + acc = acc.to(out_ptr.type.element_ty) + + offs_out_k = tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + offs_out_n = tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + + out_ptrs = ( + out_ptr + g.to(tl.int64) * K * N + offs_out_k[:, None] * N + offs_out_n[None, :] + ) + + mask = (offs_out_k[:, None] < K) & (offs_out_n[None, :] < N) + if ACCUMULATE: + # Load existing values and add to them (like beta=1 in BLAS) + old_vals = tl.load(out_ptrs, mask=mask, other=0.0) + tl.store(out_ptrs, acc + old_vals, mask=mask) + else: + # Overwrite output (like beta=0 in BLAS) + tl.store(out_ptrs, acc, mask=mask) + + # Store bias gradient (only for first N tile, sum across all M) + if COMPUTE_BIAS_GRAD and tile_n == 0: + # Keep as float32 for atomic_add (bf16/fp16 not supported for atomics) + bias_grad_ptrs = bias_grad_ptr + g.to(tl.int64) * K + offs_out_k + # Use atomic add since multiple K-tiles may write to same expert's bias + tl.atomic_add(bias_grad_ptrs, bias_acc, mask=offs_out_k < K, sem="relaxed") diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/adapter.py b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..98c224244f27445384e0c2377d73516406927536 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/adapter.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Adapt AITER's Triton grouped GEMM to MegaBlocks' ``gmm`` calling convention. + +MegaBlocks (following tgale96/grouped_gemm) uses a single ``gmm`` entry point +with ``trans_a`` / ``trans_b`` flags: + +* ``trans_a=False, trans_b=False``: a(M,K) @ b(G,K,N) -> c(M,N) +* ``trans_a=False, trans_b=True`` : a(M,K) @ b(G,N,K)^T -> c(M,N) (dgrad) +* ``trans_a=True`` : a(M,K)^T @ b(M,N) per group -> c(G,K,N) (wgrad) + +AITER exposes these as two kernels: ``gmm`` ((M,K)@(G,K,N)->(M,N), transposition +of the 3D operand inferred from strides) and ``ptgmm`` ((K,M)@(M,N)->(G,K,N), +transposition of the 2D operand inferred from strides). +""" + +import torch + +from .gmm import gmm as _aiter_gmm +from .gmm import ptgmm as _aiter_ptgmm + + +def gmm(a, b, c, batch_sizes, trans_a=False, trans_b=False): + # AITER requires group sizes to be int32 and to live on the compute device. + group_sizes = batch_sizes.to(device=a.device, dtype=torch.int32) + + # AITER asserts exact strides: gmm wants lhs/rhs row-major (a transposed + # 3D operand must be exactly column-major), tgmm wants rhs row-major and + # lhs row/column-major. Make operands contiguous first so the transposed + # views have the precise strides the kernels expect. `.contiguous()` is a + # no-op when the tensor is already contiguous. + if trans_a: + # Weight gradient: a(M,K), b(M,N) -> c(G,K,N). + # Pass a transposed so AITER sees lhs(K,M) column-major (TRANS_LHS). + _aiter_ptgmm( + a.contiguous().transpose(0, 1), + b.contiguous(), + group_sizes, + preferred_element_type=c.dtype, + existing_out=c, + ) + else: + # trans_b contracts b's last dim: pass a column-major (G,K,N) view. + rhs = b.contiguous() + if trans_b: + rhs = rhs.transpose(1, 2) + _aiter_gmm( + a.contiguous(), + rhs, + group_sizes, + preferred_element_type=c.dtype, + existing_out=c, + ) + return c diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/configs.py b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/configs.py new file mode 100644 index 0000000000000000000000000000000000000000..9a4fe5617d8100869aa76dba9b7d22c7bcab814f --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/configs.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: MIT +# Tuned GMM configs vendored from ROCm/aiter (aiter/ops/triton/configs/). +# Inlined as a Python module so packaging always includes them. + +CONFIGS = {'gfx1250': {'gmm': {'default': {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_K': 64, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}}, 'ptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}}, 'nptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}}}, 'gfx942': {'gmm': {'default': {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_K': 64, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 304, 'num_warps': 8, 'num_stages': 1}}, 'ptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 304, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'GRID_DIM': 304, 'num_warps': 8, 'num_stages': 1}}, 'nptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}}}, 'gfx950': {'gmm': {'default': {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_K': 64, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}}, 'ptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}}, 'nptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}}}} diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/gmm.py b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/gmm.py new file mode 100644 index 0000000000000000000000000000000000000000..e30c9326c6d4e4836d1303e2761ea2440a7f4750 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/gmm.py @@ -0,0 +1,567 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved. + + +# Imports. +# ------------------------------------------------------------------------------ + +# PyTorch +import torch +from torch import Tensor + +# Triton +import triton + +# AITER: GMM utility functions +from .utils.gmm_common import ( + DTYPE, + is_power_of_2, + check_input_device_dtype, + check_bias_shape_stride, + get_gmm_shape, + get_gmm_output, + get_gmm_transposition, + get_tgmm_shape, + get_tgmm_output, + get_tgmm_bias_grad, + get_tgmm_transposition, +) + +# AITER: GMM Triton kernels +from ._triton_kernels.gmm import ( + gmm_kernel, + tgmm_persistent_kernel, + tgmm_non_persistent_kernel, + get_config, +) + +# GMM PyTorch wrapper. +# ------------------------------------------------------------------------------ + + +def _gmm_grid( + N: int, + block_size_m: int, + block_size_n: int, + group_sizes: Tensor, + grid_dim: int, +) -> tuple[int]: + assert N > 0, f"N must be positive, it's {N}." + assert is_power_of_2( + block_size_m + ), f"M-dimension tile size must be a power of 2 (it's {block_size_m})." + assert is_power_of_2( + block_size_n + ), f"N-dimension tile size must be a power of 2 (it's {block_size_n})." + assert torch.all(group_sizes >= 0).item(), "All group_sizes must be non-negative." + assert grid_dim > 0, f"Grid dimension must be positive (it's {grid_dim})." + num_m_tiles = (group_sizes + block_size_m - 1) // block_size_m + assert torch.all(num_m_tiles >= 0).item(), "All num_m_tiles must be non-negative." + num_n_tiles = triton.cdiv(N, block_size_n) + assert num_n_tiles > 0, f"num_n_tiles must be positive, it's {num_n_tiles}." + num_tiles = torch.sum(num_m_tiles * num_n_tiles).item() + assert num_tiles > 0, f"num_tiles must be positive, it's {num_tiles}." + num_programs = int(min(grid_dim, num_tiles)) + assert num_programs > 0, f"num_programs must be positive, it's {num_programs}." + return (num_programs,) + + +def gmm( + lhs: Tensor, + rhs: Tensor, + group_sizes: Tensor, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, + config: dict[str, int] | None = None, + bias: Tensor | None = None, +) -> Tensor: + """ + Perform Group Matrix Multiplication (GMM): out = lhs @ rhs + bias + + lhs rows are divided into G groups. Each group of lhs rows is matrix multiplied with a plane of + rhs 3D tensor and then stored in a slice of out. In PyTorch parlance, it can be implemented as + follows for a given group g: + out[group_start:group_end, :] = lhs[group_start:group_end, :] @ rhs[g] + bias[g] + + The size of each group, and their respective start and end positions are specified by + group_sizes tensor. For instance, suppose that group_sizes = [3, 2, 4, 1]. In this particular + case we have 4 groups. The 1st group starts at 0 and ends at 2, the second group starts at 3 and + ends at 4, the third group starts at 5 and ends at 8, and the fourth and final group consists of + just the 10th (last) row of lhs. + + Parameters + ---------- + lhs : torch.Tensor + Left-hand side 2D input tensor. Shape: (M, K). + lhs data type must be torch.float16 or torch.bfloat16, and must match rhs data type. + lhs must be on the same device of rhs and group_sizes. + rhs : torch.Tensor + Right-hand side 3D input tensor. Shape: (G, K, N). + rhs data type must be torch.float16 or torch.bfloat16, and must match lhs data type. + rhs must be on the same device of lhs and group_sizes. + group_sizes : torch.Tensor + 1D input tensor describing group sizes. Shape: (G,). + group_sizes data type must be torch.int32 and all its elements must be non-negative. + group_sizes must be on the same device of lhs and rhs. + preferred_element_type : torch.dtype, optional + Desired data type for output tensor. Default is torch.bfloat16. + Supported output types are torch.float16 and torch.bfloat16. + existing_out : torch.Tensor or None, optional + Preallocated output tensor. Default is None. + If provided, results are written into this tensor. Otherwise, a new output tensor is + allocated. + If provided then it must have shape (M, N), its data type must match preferred_element_type + and it must be on the same device of other input tensors. + config : dict[str, int] or None, optional + Optional dictionary with kernel metaparameters. If absent, config will be queried from + internal tuning database. + bias : torch.Tensor or None, optional + Optional bias tensor. Shape: (G, N). + If provided, bias data type must match lhs and rhs data type, and bias must be on the same + device as other input tensors. Each group g adds bias[g] to the output. + + Returns + ------- + torch.Tensor + The computed output 2D tensor. Shape: (M, N). + Output tensor data type is given by preferred_element_type. + If existing_out is provided then existing_out is also returned. + + Implementation Notes + -------------------- + - GMM is implemented with a persistent Triton kernel. + - lhs must be row-major (lhs.stride() == (K, 1)). + - rhs can be row-major (rhs.stride() == (K * N, N, 1)) or column-major (rhs.stride() == + (K * N, 1, K)). If rhs is row-major then kernel parameter TRANS_RHS == False, this is useful + for implementing forward pass. If rhs is column-major then kernel parameter TRANS_RHS == True, + this is useful for computing the lhs derivative in the backward pass, while fusing the + transposition. + - out must be row-major (out.stride() == (N, 1)). + - bias must be row-major (bias.stride() == (N, 1)) if provided. + """ + use_bias = bias is not None + check_input_device_dtype(lhs, rhs, group_sizes, bias) + + M, K, N, G = get_gmm_shape(lhs, rhs, group_sizes) + + if use_bias: + check_bias_shape_stride(bias, G, N) + + out = get_gmm_output( + M, + N, + device=lhs.device, + preferred_element_type=preferred_element_type, + existing_out=existing_out, + ) + + trans_rhs, _ = get_gmm_transposition(lhs, rhs, out) + + if config is None: + config = get_config("gmm", M, K, N, G) + + assert all( + key in config + and isinstance(config[key], int) + and ( + is_power_of_2(config[key]) + if key.startswith("BLOCK_SIZE_") + else config[key] > 0 + ) + for key in { + "BLOCK_SIZE_M", + "BLOCK_SIZE_K", + "BLOCK_SIZE_N", + "GROUP_SIZE", + "GRID_DIM", + } + ), "Invalid GMM kernel config." + + grid = _gmm_grid( + N, + config["BLOCK_SIZE_M"], + config["BLOCK_SIZE_N"], + group_sizes, + config["GRID_DIM"], + ) + + # fmt: off + gmm_kernel[grid]( + # Tensor pointers: + lhs, rhs, group_sizes, out, bias, + # Tensor shapes: + M, K, N, G, + # Meta-parameters: + TRANS_RHS=trans_rhs, + USE_BIAS=use_bias, + **config, + ) + # fmt: on + + return out + + +# Persistent TGMM PyTorch wrapper. +# ------------------------------------------------------------------------------ + + +def _ptgmm_grid( + K: int, + N: int, + G: int, + block_size_k: int, + block_size_n: int, + grid_dim: int, +) -> tuple[int]: + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}." + assert G > 0, f"G must be positive, it's {G}." + assert is_power_of_2( + block_size_k + ), f"K-dimension tile size must be a power of 2 (it's {block_size_k})." + assert is_power_of_2( + block_size_n + ), f"N-dimension tile size must be a power of 2 (it's {block_size_n})." + assert grid_dim > 0, f"Grid dimension must be positive (it's {grid_dim})." + num_k_tiles = triton.cdiv(K, block_size_k) + assert num_k_tiles > 0, f"num_k_tiles must be positive, it's {num_k_tiles}." + num_n_tiles = triton.cdiv(N, block_size_n) + assert num_n_tiles > 0, f"num_n_tiles must be positive, it's {num_n_tiles}." + num_tiles = G * num_k_tiles * num_n_tiles + assert num_tiles > 0, f"num_tiles must be positive, it's {num_tiles}." + num_programs = min(grid_dim, num_tiles) + assert num_programs > 0, f"num_programs must be positive, it's {num_programs}." + return (num_programs,) + + +def ptgmm( + lhs: Tensor, + rhs: Tensor, + group_sizes: Tensor, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, + config: dict[str, int] | None = None, + bias_grad: Tensor | None = None, + accumulate: bool = False, +) -> Tensor: + """ + Perform a Group Matrix Multiplication (GMM) variant: out = lhs @ rhs + + lhs columns and rhs rows are divided into G groups. Each group of lhs is matrix multiplied with + the respective group of rhs and then stored in a plane of the output 3D tensor. In PyTorch + parlance, it can be implemented as follows for a given group g: + out[g] = lhs[:, group_start:group_end] @ rhs[group_start:group_end, :] + + The 't' in the operator name derives from MaxText implementation + (https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/kernels/megablox/gmm.py), + which served as the initial inspiration for this one. TGMM differs from GMM in terms of tensor + shapes. GMM does (M, K) @ (G, K, N) = (M, N) while TGMM does (K, M) @ (M, N) = (G, K, N). + + The 'p' in the operator name means that it is implemented with a persistent kernel. There is + also the non-persistent variation, which is implemented with a regular kernel. Please take a + look at nptgmm operator. Both ptgmm and nptgmm implement the same computation, choosing one or + the other is a matter of performance for the target workload. + + Parameters + ---------- + lhs : torch.Tensor + Left-hand side 2D input tensor. Shape: (K, M). + lhs data type must be torch.float16 or torch.bfloat16, and must match rhs data type. + lhs must be on the same device of rhs and group_sizes. + rhs : torch.Tensor + Right-hand side 2D input tensor. Shape: (M, N). + rhs data type must be torch.float16 or torch.bfloat16, and must match lhs data type. + rhs must be on the same device of lhs and group_sizes. + group_sizes : torch.Tensor + 1D input tensor describing group sizes. Shape: (G,). + group_sizes data type must be torch.int32 and all its elements must be non-negative. + group_sizes must be on the same device of lhs and rhs. + preferred_element_type : torch.dtype, optional + Desired data type for output tensor. Default is torch.bfloat16. + Supported output types are torch.float16 and torch.bfloat16. + existing_out : torch.Tensor or None, optional + Preallocated output tensor. Default is None. + If provided, results are written into this tensor. Otherwise, a new output tensor is + allocated. + If provided then it must have shape (G, K, N), its data type must match + preferred_element_type and it must be on the same device of other input tensors. + config : dict[str, int] or None, optional + Optional dictionary with kernel metaparameters. If absent, config will be queried from + internal tuning database. + bias_grad : torch.Tensor or None, optional + Optional bias gradient output tensor. Shape: (G, K). + If provided, the kernel will compute the bias gradient and write it to this tensor. + bias_grad must be torch.float32 (kernel uses atomic_add which requires float32), + accumulate : bool, optional + Whether to accumulate into existing output tensor values. Default is False. + If False, output will be overwritten with fresh computation. + If True, results will be added to existing output tensor values. + + Returns + ------- + torch.Tensor + The computed output 3D tensor. Shape: (G, K, N). + Output tensor data type is given by preferred_element_type. + If existing_out is provided then existing_out is also returned. + + Implementation Notes + -------------------- + - PTGMM is implemented with a persistent Triton kernel. + - lhs can be row-major (lhs.stride() == (M, 1)) or column-major (lhs.stride() == (1, K)). If lhs + is row-major then kernel parameter TRANS_LHS == False. If lhs is column-major then kernel + parameter TRANS_LHS == True, this is useful for computing the rhs derivative in the backward + pass, while fusing the transposition. + - rhs must be row-major (rhs.stride() == (N, 1)). + - out must be row-major (out.stride() == (K * N, N, 1)). + """ + check_input_device_dtype(lhs, rhs, group_sizes) + + M, K, N, G = get_tgmm_shape(lhs, rhs, group_sizes) + + out = get_tgmm_output( + K, + N, + G, + device=lhs.device, + preferred_element_type=preferred_element_type, + existing_out=existing_out, + ) + + trans_lhs, _ = get_tgmm_transposition(lhs, rhs, out) + + if config is None: + config = get_config("ptgmm", M, K, N, G, accumulate) + + assert all( + key in config + and isinstance(config[key], int) + and ( + is_power_of_2(config[key]) + if key.startswith("BLOCK_SIZE_") + else config[key] > 0 + ) + for key in { + "BLOCK_SIZE_M", + "BLOCK_SIZE_K", + "BLOCK_SIZE_N", + "GROUP_SIZE", + "GRID_DIM", + } + ), "Invalid PTGMM kernel config." + + # Bias gradient handling. + # ----------------------- + # Get or validate bias gradient tensor. + compute_bias_grad = bias_grad is not None + bias_grad_ptr = get_tgmm_bias_grad( + K, + G, + device=lhs.device, + existing_bias_grad=bias_grad, + ) + + grid = _ptgmm_grid( + K, + N, + G, + config["BLOCK_SIZE_K"], + config["BLOCK_SIZE_N"], + config["GRID_DIM"], + ) + + # fmt: off + tgmm_persistent_kernel[grid]( + # Tensor pointers: + lhs, rhs, group_sizes, out, bias_grad_ptr, + # Tensor shapes: + M, K, N, G, + # Meta-parameters: + TRANS_LHS=trans_lhs, + COMPUTE_BIAS_GRAD=compute_bias_grad, + ACCUMULATE=accumulate, + **config, + ) + # fmt: on + + return out + + +# Regular non-persistent TGMM PyTorch wrapper. +# ------------------------------------------------------------------------------ + + +def _nptgmm_grid( + K: int, + N: int, + G: int, + block_size_k: int, + block_size_n: int, +) -> tuple[int, int]: + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}." + assert G > 0, f"G must be positive, it's {G}." + assert is_power_of_2( + block_size_k + ), f"K-dimension tile size must be a power of 2 (it's {block_size_k})." + assert is_power_of_2( + block_size_n + ), f"N-dimension tile size must be a power of 2 (it's {block_size_n})." + num_k_tiles = triton.cdiv(K, block_size_k) + assert num_k_tiles > 0, f"num_k_tiles must be positive, it's {num_k_tiles}." + num_n_tiles = triton.cdiv(N, block_size_n) + assert num_n_tiles > 0, f"num_n_tiles must be positive, it's {num_n_tiles}." + num_tiles_per_mm = num_k_tiles * num_n_tiles + assert ( + num_tiles_per_mm > 0 + ), f"num_tiles_per_mm must be positive, it's {num_tiles_per_mm}." + return (G, num_tiles_per_mm) + + +def nptgmm( + lhs: Tensor, + rhs: Tensor, + group_sizes: Tensor, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, + config: dict[str, int] | None = None, + bias_grad: Tensor | None = None, + accumulate: bool = False, +) -> Tensor: + """ + Perform a Group Matrix Multiplication (GMM) variant: out = lhs @ rhs + + lhs columns and rhs rows are divided into G groups. Each group of lhs is matrix multiplied with + the respective group of rhs and then stored in a plane of the output 3D tensor. In PyTorch + parlance, it can be implemented as follows for a given group g: + out[g] = lhs[:, group_start:group_end] @ rhs[group_start:group_end, :] + + The 't' in the operator name derives from MaxText implementation + (https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/kernels/megablox/gmm.py), + which served as the initial inspiration for this one. TGMM differs from GMM in terms of tensor + shapes. GMM does (M, K) @ (G, K, N) = (M, N) while TGMM does (K, M) @ (M, N) = (G, K, N). + + The 'np' in the operator name means that it is implemented with a non-persistent, i.e. regular + kernel. There is also the persistent variation, which is implemented with a persistent kernel. + Please take a look at ptgmm operator. Both nptgmm and ptgmm implement the same computation, + choosing one or the other is a matter of performance for the target workload. + + Parameters + ---------- + lhs : torch.Tensor + Left-hand side 2D input tensor. Shape: (K, M). + lhs data type must be torch.float16 or torch.bfloat16, and must match rhs data type. + lhs must be on the same device of rhs and group_sizes. + rhs : torch.Tensor + Right-hand side 2D input tensor. Shape: (M, N). + rhs data type must be torch.float16 or torch.bfloat16, and must match lhs data type. + rhs must be on the same device of lhs and group_sizes. + group_sizes : torch.Tensor + 1D input tensor describing group sizes. Shape: (G,). + group_sizes data type must be torch.int32 and all its elements must be non-negative. + group_sizes must be on the same device of lhs and rhs. + preferred_element_type : torch.dtype, optional + Desired data type for output tensor. Default is torch.bfloat16. + Supported output types are torch.float16 and torch.bfloat16. + existing_out : torch.Tensor or None, optional + Preallocated output tensor. Default is None. + If provided, results are written into this tensor. Otherwise, a new output tensor is + allocated. + If provided then it must have shape (G, K, N), its data type must match + preferred_element_type and it must be on the same device of other input tensors. + config : dict[str, int] or None, optional + Optional dictionary with kernel metaparameters. If absent, config will be queried from + internal tuning database. + bias_grad : torch.Tensor or None, optional + Optional bias gradient output tensor. Shape: (G, K). + If provided, the kernel will compute the bias gradient and write it to this tensor. + bias_grad must be torch.float32 (kernel uses atomic_add which requires float32), + accumulate : bool, optional + Whether to accumulate into existing output tensor values. Default is False. + If False, output will be overwritten with fresh computation. + If True, results will be added to existing output tensor values. + + Returns + ------- + torch.Tensor + The computed output 3D tensor. Shape: (G, K, N). + Output tensor data type is given by preferred_element_type. + If existing_out is provided then existing_out is also returned. + + Implementation Notes + -------------------- + - NPTGMM is implemented with a non-persistent regular Triton kernel. + - lhs can be row-major (lhs.stride() == (M, 1)) or column-major (lhs.stride() == (1, K)). If lhs + is row-major then kernel parameter TRANS_LHS == False. If lhs is column-major then kernel + parameter TRANS_LHS == True, this is useful for computing the rhs derivative in the backward + pass, while fusing the transposition. + - rhs must be row-major (rhs.stride() == (N, 1)). + - out must be row-major (out.stride() == (K * N, N, 1)). + """ + check_input_device_dtype(lhs, rhs, group_sizes) + + M, K, N, G = get_tgmm_shape(lhs, rhs, group_sizes) + + out = get_tgmm_output( + K, + N, + G, + device=lhs.device, + preferred_element_type=preferred_element_type, + existing_out=existing_out, + ) + + trans_lhs, _ = get_tgmm_transposition(lhs, rhs, out) + + # Bias gradient handling. + # ----------------------- + # Get or validate bias gradient tensor. + compute_bias_grad = bias_grad is not None + bias_grad_ptr = get_tgmm_bias_grad( + K, + G, + device=lhs.device, + existing_bias_grad=bias_grad, + ) + + if config is None: + config = get_config("nptgmm", M, K, N, G, accumulate) + + assert all( + key in config + and isinstance(config[key], int) + and ( + is_power_of_2(config[key]) + if key.startswith("BLOCK_SIZE_") + else config[key] > 0 + ) + for key in { + "BLOCK_SIZE_M", + "BLOCK_SIZE_K", + "BLOCK_SIZE_N", + "GROUP_SIZE", + } + ), "Invalid NPTGMM kernel config." + + grid = _nptgmm_grid( + K, + N, + G, + config["BLOCK_SIZE_K"], + config["BLOCK_SIZE_N"], + ) + + # fmt: off + tgmm_non_persistent_kernel[grid]( + # Tensor pointers: + lhs, rhs, group_sizes, out, bias_grad_ptr, + # Tensor shapes: + M, K, N, G, + # Meta-parameters: + TRANS_LHS=trans_lhs, + COMPUTE_BIAS_GRAD=compute_bias_grad, + ACCUMULATE=accumulate, + **config, + ) + # fmt: on + + return out diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/_triton/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/_triton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/_triton/arch_info.py b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/_triton/arch_info.py new file mode 100644 index 0000000000000000000000000000000000000000..3f6c88581a64044518125623f116082c53bd5474 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/_triton/arch_info.py @@ -0,0 +1,46 @@ +import triton + +# Detect the GPU arch lazily: querying the triton driver at import time fails +# in headless environments (e.g. the kernel-builder ABI check sandbox has no +# GPU), and the original JAX fallback pulled in an unrelated runtime dep. The +# arch is only actually needed when a GMM kernel is dispatched, so resolve and +# cache on first call. +_CACHED_ARCH = None + + +def get_arch(): + global _CACHED_ARCH + if _CACHED_ARCH is not None: + return _CACHED_ARCH + try: + _CACHED_ARCH = triton.runtime.driver.active.get_current_target().arch + except RuntimeError: + try: + from jax._src.lib import gpu_triton as triton_kernel_call_lib + _CACHED_ARCH = triton_kernel_call_lib.get_arch_details("0").split(":")[0] + except ImportError as e: + raise RuntimeError( + "Cannot determine GPU arch: triton driver is inactive and " + "JAX is not available. A GPU is required for grouped GEMM." + ) from e + return _CACHED_ARCH + + +def is_gluon_avail(): + return get_arch() in ("gfx950", "gfx1250") + + +def is_fp4_avail(): + return get_arch() in ("gfx950", "gfx1250") + + +def is_fp8_avail(): + return get_arch() in ("gfx942", "gfx950", "gfx1250", "gfx1200", "gfx1201") + + +def is_mx_scale_preshuffling_avail(): + return get_arch() in ("gfx950", "gfx1250") + + +def is_tdm_avail(): + return get_arch() in ("gfx1250",) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/_triton/pid_preprocessing.py b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/_triton/pid_preprocessing.py new file mode 100644 index 0000000000000000000000000000000000000000..99792bb3ba2fab8fff223bba733ced1eb6e6df53 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/_triton/pid_preprocessing.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: MIT + +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import triton +import triton.language as tl + + +@triton.jit +def remap_xcd_chunked( + pid, GRID_MN, NUM_XCDS: tl.constexpr = 8, CHUNK_SIZE: tl.constexpr = 2 +): + # Compute current XCD and local PID + xcd = pid % NUM_XCDS + # distribute the modulo pids in round robin + if pid > (GRID_MN // (NUM_XCDS * CHUNK_SIZE)) * (NUM_XCDS * CHUNK_SIZE): + return pid + local_pid = pid // NUM_XCDS + # Calculate chunk index and position within chunk + chunk_idx = local_pid // CHUNK_SIZE + pos_in_chunk = local_pid % CHUNK_SIZE + # Calculate new PID + new_pid = chunk_idx * NUM_XCDS * CHUNK_SIZE + xcd * CHUNK_SIZE + pos_in_chunk + return new_pid + + +@triton.jit +def remap_xcd(pid, GRID_MN, NUM_XCDS: tl.constexpr = 8): + ## pid remapping on xcds + # Number of pids per XCD in the new arrangement + pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS + # When GRID_MN cannot divide NUM_XCDS, some xcds will have + # pids_per_xcd pids, the other will have pids_per_xcd - 1 pids. + # We calculate the number of xcds that have pids_per_xcd pids as + # tall_xcds + tall_xcds = GRID_MN % NUM_XCDS + tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds + # Compute current XCD and local pid within the XCD + xcd = pid % NUM_XCDS + local_pid = pid // NUM_XCDS + # Calculate new pid based on the new grouping + # Note that we need to consider the following two cases: + # 1. the current pid is on a tall xcd + # 2. the current pid is on a short xcd + if xcd < tall_xcds: + pid = xcd * pids_per_xcd + local_pid + else: + pid = ( + tall_xcds * pids_per_xcd + + (xcd - tall_xcds) * (pids_per_xcd - 1) + + local_pid + ) + + return pid + + +@triton.jit +def pid_grid(pid: int, num_pid_m: int, num_pid_n: int, GROUP_SIZE_M: tl.constexpr = 1): + """ + Maps 1D pid to 2D grid coords (pid_m, pid_n). + + Args: + - pid: 1D pid + - num_pid_m: grid m size + - num_pid_n: grid n size + - GROUP_SIZE_M: tl.constexpr: default is 1 + """ + if GROUP_SIZE_M == 1: + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + else: + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + tl.assume(group_size_m >= 0) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + return pid_m, pid_n + + +@triton.jit +def pid_grid_3d(pid: int, num_pid_m: int, num_pid_n: int, num_pid_k): + """ + Maps 1D pid to 3D grid coords (pid_m, pid_n, pid_k). + Args: + - pid: 1D pid + - num_pid_m: grid m size + - num_pid_n: grid n size + - num_pid_k: grid k size + + Returns: + - pid_m, pid_n, pid_k: 3D grid coordinates + """ + pid_m = pid % num_pid_m + pid_n = (pid // num_pid_m) % num_pid_n + pid_k = pid // (num_pid_m * num_pid_n) % num_pid_k + + return pid_m, pid_n, pid_k diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/gmm_common.py b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/gmm_common.py new file mode 100644 index 0000000000000000000000000000000000000000..153dee65b50ab5f833262481889d2184d1ca639f --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/gmm_common.py @@ -0,0 +1,752 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved. + +# Imports. +# ------------------------------------------------------------------------------ + +# PyTorch +import torch +from torch import Tensor + +# AITER: logging +from .logger import AiterTritonLogger + +_LOGGER: AiterTritonLogger = AiterTritonLogger() + + +# Supported data types. +# ------------------------------------------------------------------------------ + +# Supported data types, as strings. +SUPPORTED_DTYPES_STR: set[str] = {"fp16", "bf16"} + + +# Convert string data type to PyTorch data type. +def dtype_from_str(dtype_str: str) -> torch.dtype: + dtype_str = dtype_str.strip().lower() + dtype_str = dtype_str[1:] if dtype_str[0] in {"i", "o"} else dtype_str + assert ( + dtype_str in SUPPORTED_DTYPES_STR + ), "String data type isn't in set of supported string data types." + return {"fp16": torch.float16, "bf16": torch.bfloat16}[dtype_str] + + +# Supported data types, as PyTorch types. +SUPPORTED_DTYPES: set[torch.dtype] = { + dtype_from_str(dtype_str) for dtype_str in SUPPORTED_DTYPES_STR +} + + +# Convert PyTorch data type to string data type. +def str_from_dtype(dtype: torch.dtype) -> str: + assert ( + dtype in SUPPORTED_DTYPES + ), "PyTorch data type isn't in set of supported PyTorch data types." + return {torch.float16: "fp16", torch.bfloat16: "bf16"}[dtype] + + +# Default data type, as string. +DTYPE_STR: str = "bf16" +assert ( + DTYPE_STR in SUPPORTED_DTYPES_STR +), "Default string data type isn't in set of supported string data types." + + +# Default data type, as PyTorch type. +DTYPE: torch.dtype = dtype_from_str(DTYPE_STR) + + +# Other defaults. +# ------------------------------------------------------------------------------ + +# Default device. +DEVICE: torch.device | str = "cuda" + +# Default RNG seed for input generation. +RNG_SEED: int = 0 + +# Default number of group sizes. +NUM_GROUP_SIZES: int = 1 + +# Default transposition (NN). +TRANS_LHS: bool = False +TRANS_RHS: bool = False + + +# Parameter checking functions. +# ------------------------------------------------------------------------------ + + +def is_power_of_2(x: int) -> bool: + return (x > 0) and (x & (x - 1) == 0) + + +def check_input_device_dtype( + lhs: Tensor, rhs: Tensor, group_sizes: Tensor, bias: Tensor | None = None +) -> None: + assert ( + lhs.device == rhs.device == group_sizes.device + ), f"All input tensors must be in the same device (lhs = {lhs.device}, rhs = {rhs.device}, group_sizes = {group_sizes.device})." + assert ( + lhs.dtype == rhs.dtype + ), f"lhs and rhs types must match (lhs = {lhs.dtype}, rhs = {rhs.dtype})." + assert group_sizes.dtype == torch.int32, "group_sizes type must be int32." + + if bias is not None: + assert ( + bias.device == lhs.device + ), f"bias must be on the same device as lhs (bias = {bias.device}, lhs = {lhs.device})." + assert ( + bias.dtype == lhs.dtype + ), f"bias dtype must match lhs dtype (bias = {bias.dtype}, lhs = {lhs.dtype})." + + +def check_bias_shape_stride(bias: Tensor, G: int, N: int) -> None: + assert bias.shape == ( + G, + N, + ), f"bias must have shape (G, N) = ({G}, {N}), got {bias.shape}." + assert bias.stride() == (N, 1), "bias must be row-major (bias.stride() == (N, 1))." + + +# Generation of group sizes. +# ------------------------------------------------------------------------------ + + +# Probabilities for generating random group sizes. +UNUSED_TOKENS_PROB: float = 0.0 +UNUSED_EXPERTS_PROB: float = 0.1 + + +def gen_uniform_group_sizes( + M: int, + G: int, + device: torch.device | str = DEVICE, +) -> Tensor: + assert M >= 0, f"Number of tokens M must be non-negative (it's {M})." + assert G > 0, f"Number of experts G must be positive (it's {G})." + + base = M // G + remainder = M % G + group_sizes = torch.full((G,), base, dtype=torch.int32, device=device) + if remainder > 0: + group_sizes[:remainder] += 1 + + assert ( + len(group_sizes) == G + ), f"Group sizes don't have {G} elements (it's {len(group_sizes)})." + assert torch.all(group_sizes >= 0).item(), "All group sizes must be non-negative." + assert ( + torch.sum(group_sizes).item() == M + ), f"Group sizes don't add up to total tokens {M}." + assert group_sizes.dtype == torch.int32, "Group sizes must be int32." + + return group_sizes + + +def gen_group_sizes( + M: int, + G: int, + device: torch.device | str = DEVICE, + rng_seed: int | None = RNG_SEED, + unused_tokens_prob: float = UNUSED_TOKENS_PROB, + unused_experts_prob: float = UNUSED_EXPERTS_PROB, +) -> Tensor: + assert M >= 0, f"Number of tokens M must be non-negative (it's {M})." + assert G > 0, f"Number of experts G must be positive (it's {G})." + assert ( + 0 <= unused_tokens_prob <= 1 + ), f"Probability of unused tokens must be in [0, 1] interval (it's {unused_tokens_prob})." + assert ( + 0 <= unused_experts_prob <= 1 + ), f"Probability of unused experts must be in [0, 1] interval (it's {unused_experts_prob})." + + if rng_seed is not None: + torch.manual_seed(rng_seed) + + if unused_tokens_prob > 0: + # Optionally drop tokens to simulate routing sparsity, some tokens may not be routed. + num_unused_tokens = M + while num_unused_tokens == M: + num_unused_tokens = int( + torch.binomial( + torch.tensor(float(M), device=device), + torch.tensor(unused_tokens_prob, device=device), + ).item() + ) + else: + num_unused_tokens = 0 + num_used_tokens = M - num_unused_tokens + assert ( + num_unused_tokens >= 0 + ), f"Number of unused tokens must be non-negative (it's {num_unused_tokens})." + assert ( + num_used_tokens > 0 + ), f"Number of used tokens must be positive (it's {num_used_tokens})." + assert ( + num_used_tokens + num_unused_tokens == M + ), f"Unused + used tokens don't add up total tokens ({num_used_tokens} + {num_unused_tokens} != {M})." + + if num_unused_tokens > 0: + _LOGGER.debug( + f"Group sizes generation: dropped {num_unused_tokens} token{'s' if num_unused_tokens > 1 else ''}.", + ) + + if unused_experts_prob > 0: + # Some experts may have zero tokens assigned to them. + num_used_experts = 0 + while num_used_experts == 0: + used_experts = torch.nonzero( + torch.rand((G,), device=device) >= unused_experts_prob + ).squeeze() + num_used_experts = used_experts.numel() + else: + used_experts = torch.arange(0, G, device=device) + num_used_experts = G + num_unused_experts = G - num_used_experts + assert ( + num_unused_experts >= 0 + ), f"Number of unused experts must be non-negative (it's {num_unused_experts})." + assert ( + num_used_experts >= 1 + ), f"At least one expert must be used (it's {num_used_experts})." + assert ( + num_unused_experts + num_used_experts == G + ), f"Unused + used experts don't add up total experts ({num_unused_experts} + {num_used_experts} != {G})." + + if num_unused_experts > 0: + _LOGGER.debug( + f"Group sizes generation: dropped {num_unused_experts} expert{'s' if num_unused_experts > 1 else ''}.", + ) + + group_sizes = torch.bincount( + used_experts[ + torch.randint(low=0, high=num_used_experts, size=(num_used_tokens,)) + ], + minlength=G, + ).to(torch.int32) + + assert ( + len(group_sizes) == G + ), f"Group sizes don't have {G} elements (it's {len(group_sizes)})." + assert torch.all(group_sizes >= 0).item(), "All group sizes must be non-negative." + assert ( + torch.sum(group_sizes).item() == num_used_tokens + ), f"Group sizes don't add up to used tokens {num_used_tokens}." + assert group_sizes.dtype == torch.int32, "Group sizes must be int32." + + return group_sizes + + +def gen_multiple_group_sizes( + num_group_sizes: int, + M: int, + G: int, + device: torch.device | str = DEVICE, + rng_seed: int | None = RNG_SEED, + unused_tokens_prob: float = UNUSED_TOKENS_PROB, + unused_experts_prob: float = UNUSED_EXPERTS_PROB, + group_sizes_0: Tensor | None = None, +) -> list[Tensor]: + assert ( + num_group_sizes > 0 + ), f"Number of group sizes to be generated must be positive, it's {num_group_sizes}." + multiple_group_sizes = [ + gen_group_sizes( + M, + G, + device=device, + rng_seed=rng_seed if g == 0 else None, + unused_tokens_prob=unused_tokens_prob, + unused_experts_prob=unused_experts_prob, + ) + for g in range( + num_group_sizes if group_sizes_0 is None else num_group_sizes - 1 + ) + ] + if group_sizes_0 is not None: + multiple_group_sizes.insert(0, group_sizes_0) + assert ( + len(multiple_group_sizes) == num_group_sizes + ), f"Expecting {num_group_sizes} distinct group sizes (it's {len(multiple_group_sizes)})." + return multiple_group_sizes + + +# GMM helpers: tensor generation. +# ------------------------------------------------------------------------------ + + +def gen_gmm_input( + M: int, + K: int, + N: int, + G: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, + trans_rhs: bool = TRANS_RHS, + rng_seed: int | None = RNG_SEED, + unif_group_sizes: bool = False, +) -> tuple[Tensor, Tensor, Tensor]: + assert M > 0, f"Number of lhs rows M must be positive (M = {M})." + assert K > 0, f"Number of lhs columns / rhs rows K must be positive (K = {K})." + assert N > 0, f"Number of rhs columns N must be positive (N = {N})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + if rng_seed is not None: + torch.manual_seed(rng_seed) + + lhs = torch.randn((M, K), dtype=torch.float32, device=device) + lhs = lhs.to(preferred_element_type) + + if trans_rhs: + rhs = torch.randn((G, N, K), dtype=torch.float32, device=device).permute( + 0, 2, 1 + ) + else: + rhs = torch.randn((G, K, N), dtype=torch.float32, device=device) + rhs = rhs.to(preferred_element_type) + + group_sizes = ( + gen_uniform_group_sizes(M, G, device=device) + if unif_group_sizes + else gen_group_sizes(M, G, device=device, rng_seed=None) + ) + + return lhs, rhs, group_sizes + + +def gen_gmm_output( + M: int, + N: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, +) -> Tensor: + assert M > 0, f"Number of out rows M must be positive (M = {M})." + assert N > 0, f"Number of out columns N must be positive (N = {N})." + + out = torch.empty((M, N), dtype=preferred_element_type, device=device) + + return out + + +def gen_gmm_tensors( + M: int, + K: int, + N: int, + G: int, + num_group_sizes: int, + device: torch.device | str = DEVICE, + input_type: torch.dtype = DTYPE, + output_type: torch.dtype = DTYPE, + trans_lhs: bool = False, + trans_rhs: bool = TRANS_RHS, + rng_seed: int | None = RNG_SEED, + unif_group_sizes: bool = False, + use_bias: bool = False, +) -> tuple[Tensor, Tensor, list[Tensor], Tensor, Tensor | None]: + lhs, rhs, group_sizes_0 = gen_gmm_input( + M, + K, + N, + G, + device=device, + preferred_element_type=input_type, + trans_rhs=trans_rhs, + rng_seed=rng_seed, + unif_group_sizes=unif_group_sizes, + ) + multiple_group_sizes = gen_multiple_group_sizes( + num_group_sizes, M, G, device=device, rng_seed=None, group_sizes_0=group_sizes_0 + ) + out = gen_gmm_output(M, N, device=device, preferred_element_type=output_type) + bias = None + if use_bias: + torch.manual_seed(rng_seed + 1000) # Different seed for bias + bias = torch.randn(G, N, dtype=input_type, device=device) + + return lhs, rhs, multiple_group_sizes, out, bias + + +# GMM helpers: get information from tensors. +# ------------------------------------------------------------------------------ + + +def get_gmm_shape( + lhs: Tensor, rhs: Tensor, group_sizes: Tensor +) -> tuple[int, int, int, int]: + assert lhs.dim() == 2, f"lhs must have 2 dimensions (it's {lhs.dim()})." + assert rhs.dim() == 3, f"rhs must have 3 dimensions (it's {rhs.dim()})." + assert ( + group_sizes.dim() == 1 + ), f"group_sizes must have 1 dimension (it's {group_sizes.dim()})." + + M, lhs_k = lhs.shape + rhs_g, rhs_k, N = rhs.shape + group_sizes_g = group_sizes.shape[0] + + assert ( + lhs_k == rhs_k + ), f"K dimension of lhs and rhs don't match (lhs = {lhs_k}, rhs = {rhs_k})." + K = lhs_k + assert ( + rhs_g == group_sizes_g + ), f"G dimension of rhs and group_sizes don't match (rhs = {rhs_g}, group_sizes = {group_sizes_g})." + G = rhs_g + + assert M > 0, f"M must be positive, it's {M}." + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}" + assert G > 0, f"G must be positive, it's {G}" + + return M, K, N, G + + +def get_gmm_output( + M: int, + N: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, +) -> Tensor: + assert M > 0, f"Number of out rows M must be positive (M = {M})." + assert N > 0, f"Number of out columns N must be positive (N = {N})." + + if existing_out is not None: + assert ( + existing_out.device == device + ), f"Existing output device and provided device don't match (existing = {existing_out.device}, provided = {device})." + assert ( + existing_out.dtype == preferred_element_type + ), f"Existing output type and preferred output type don't match (existing = {existing_out.dtype}, preferred = {preferred_element_type})." + assert existing_out.shape == ( + M, + N, + ), f"Existing output shape and GMM shape don't match (existing = {tuple(existing_out.shape)}, provided = {(M, N)})." + return existing_out + + return gen_gmm_output( + M, + N, + device=device, + preferred_element_type=preferred_element_type, + ) + + +def get_gmm_transposition(lhs: Tensor, rhs: Tensor, out: Tensor) -> tuple[bool, int]: + assert lhs.dim() == 2, f"lhs must have 2 dimensions (it's {lhs.dim()})." + assert rhs.dim() == 3, f"rhs must have 3 dimensions (it's {rhs.dim()})." + assert out.dim() == 2, f"out must have 2 dimensions (it's {out.dim()})." + + lhs_m, lhs_k = lhs.shape + G, rhs_k, rhs_n = rhs.shape + out_m, out_n = out.shape + + assert ( + lhs_m == out_m + ), f"M dimension of lhs and out don't match (lhs = {lhs_m}, rhs = {out_m})." + M = lhs_m + assert ( + lhs_k == rhs_k + ), f"K dimension of lhs and rhs don't match (lhs = {lhs_k}, rhs = {rhs_k})." + K = lhs_k + assert ( + rhs_n == out_n + ), f"N dimension of rhs and out don't match (lhs = {rhs_n}, rhs = {out_n})." + N = rhs_n + + assert M > 0, f"M must be positive, it's {M}." + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}" + assert G > 0, f"G must be positive, it's {G}" + + is_lhs_row_major = lhs.stride() == (K, 1) + assert is_lhs_row_major, "lhs must be row-major." + is_rhs_row_major = rhs.stride() == (K * N, N, 1) + is_rhs_col_major = rhs.stride() == (K * N, 1, K) + assert ( + is_rhs_row_major != is_rhs_col_major + ), "rhs must be row-major or column-major." + is_out_row_major = out.stride() == (N, 1) + assert is_out_row_major, "out must be row-major." + + # Get rhs leading dimension according to transposition configuration. + ld_rhs = N if is_rhs_row_major else K + + return is_rhs_col_major, ld_rhs + + +# TGMM helpers: tensor generation. +# ------------------------------------------------------------------------------ + + +def gen_tgmm_input( + M: int, + K: int, + N: int, + G: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, + trans_lhs: bool = TRANS_LHS, + rng_seed: int | None = RNG_SEED, + unif_group_sizes: bool = False, +) -> tuple[Tensor, Tensor, Tensor]: + assert K > 0, f"Number of lhs rows K must be positive (M = {K})." + assert M > 0, f"Number of lhs columns / rhs rows M must be positive (K = {M})." + assert N > 0, f"Number of rhs columns N must be positive (N = {N})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + if rng_seed is not None: + torch.manual_seed(rng_seed) + + if trans_lhs: + lhs = torch.randn((M, K), dtype=torch.float32, device=device).T + else: + lhs = torch.randn((K, M), dtype=torch.float32, device=device) + lhs = lhs.to(preferred_element_type) + + rhs = torch.randn((M, N), dtype=torch.float32, device=device) + rhs = rhs.to(preferred_element_type) + + group_sizes = ( + gen_uniform_group_sizes(M, G, device=device) + if unif_group_sizes + else gen_group_sizes(M, G, device=device, rng_seed=None) + ) + + return lhs, rhs, group_sizes + + +def gen_tgmm_output( + K: int, + N: int, + G: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, +) -> Tensor: + assert K > 0, f"Number of out rows K must be positive (K = {K})." + assert N > 0, f"Number of out columns N must be positive (N = {N})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + out = torch.empty((G, K, N), dtype=preferred_element_type, device=device) + + return out + + +def gen_tgmm_bias_grad( + K: int, + G: int, + device: torch.device | str = DEVICE, + with_bias_grad: bool = False, +) -> Tensor: + if with_bias_grad: + assert K > 0, f"Number of bias_grad rows K must be positive (K = {K})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + return torch.empty((G, K), device=device, dtype=torch.float32) + else: + # Return dummy pointer when bias_grad is not needed. + # Must be float32 because atomic_add does not support bf16/fp16, + # and Triton validates the pointer dtype even in dead branches. + return torch.tensor([], device=device, dtype=torch.float32) + + +def gen_tgmm_tensors( + M: int, + K: int, + N: int, + G: int, + num_group_sizes: int, + device: torch.device | str = DEVICE, + input_type: torch.dtype = DTYPE, + output_type: torch.dtype = DTYPE, + trans_lhs: bool = TRANS_LHS, + trans_rhs: bool = False, + rng_seed: int | None = RNG_SEED, + unif_group_sizes: bool = False, + use_bias: bool = False, +) -> tuple[Tensor, Tensor, list[Tensor], Tensor, Tensor | None]: + lhs, rhs, group_sizes_0 = gen_tgmm_input( + M, + K, + N, + G, + device=device, + preferred_element_type=input_type, + trans_lhs=trans_lhs, + rng_seed=rng_seed, + unif_group_sizes=unif_group_sizes, + ) + multiple_group_sizes = gen_multiple_group_sizes( + num_group_sizes, M, G, device=device, rng_seed=None, group_sizes_0=group_sizes_0 + ) + out = gen_tgmm_output(K, N, G, device=device, preferred_element_type=output_type) + if use_bias: + bias_grad = gen_tgmm_bias_grad(K, G, device=device, with_bias_grad=True) + else: + bias_grad = None + return lhs, rhs, multiple_group_sizes, out, bias_grad + + +# TGMM helpers: get information from tensors. +# ------------------------------------------------------------------------------ + + +def get_tgmm_shape( + lhs: Tensor, rhs: Tensor, group_sizes: Tensor +) -> tuple[int, int, int, int]: + assert lhs.dim() == 2, f"lhs must have 2 dimensions (it's {lhs.dim()})." + assert rhs.dim() == 2, f"rhs must have 2 dimensions (it's {rhs.dim()})." + assert ( + group_sizes.dim() == 1 + ), f"group_sizes must have 1 dimension (it's {group_sizes.dim()})." + + K, lhs_m = lhs.shape + rhs_m, N = rhs.shape + G = group_sizes.shape[0] + + assert ( + lhs_m == rhs_m + ), f"M dimension of lhs and rhs don't match (lhs = {lhs_m}, rhs = {rhs_m})." + M = lhs_m + + assert M > 0, f"M must be positive, it's {M}." + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}" + assert G > 0, f"G must be positive, it's {G}" + + return M, K, N, G + + +def get_tgmm_output( + K: int, + N: int, + G: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, +) -> Tensor: + assert K > 0, f"Number of out rows K must be positive (K = {K})." + assert N > 0, f"Number of out columns N must be positive (N = {N})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + if existing_out is not None: + assert ( + existing_out.device == device + ), f"Existing output device and provided device don't match (existing = {existing_out.device}, provided = {device})." + assert ( + existing_out.dtype == preferred_element_type + ), f"Existing output type and preferred output type don't match (existing = {existing_out.dtype}, preferred = {preferred_element_type})." + assert existing_out.shape == ( + G, + K, + N, + ), f"Existing output shape and GMM shape don't match (existing = {tuple(existing_out.shape)}, provided = {(G, K, N)})." + return existing_out + + return gen_tgmm_output( + K, + N, + G, + device=device, + preferred_element_type=preferred_element_type, + ) + + +def get_tgmm_bias_grad( + K: int, + G: int, + device: torch.device | str = DEVICE, + existing_bias_grad: Tensor | None = None, +) -> Tensor: + """ + Get or validate bias gradient tensor for TGMM. + + If existing_bias_grad is provided, validates its shape, device, dtype, and stride, + and always zeros it before returning (since the kernel uses atomic_add). + If existing_bias_grad is None, returns a dummy tensor (for use when COMPUTE_BIAS_GRAD=False). + Parameters + ---------- + K : int + Number of rows in the bias gradient tensor. + G : int + Number of groups. + device : torch.device or str + Device for the tensor. + existing_bias_grad : torch.Tensor or None + Existing bias gradient tensor to validate and use. + Returns + ------- + torch.Tensor + Valid bias gradient tensor or dummy tensor. + """ + assert K > 0, f"Number of bias_grad rows K must be positive (K = {K})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + if existing_bias_grad is not None: + # Validate existing bias_grad tensor. + expected_shape = (G, K) + assert ( + tuple(existing_bias_grad.shape) == expected_shape + ), f"bias_grad must have shape {expected_shape}, got {tuple(existing_bias_grad.shape)}." + assert ( + existing_bias_grad.device == device + ), f"bias_grad must be on the same device (bias_grad = {existing_bias_grad.device}, device = {device})." + assert ( + existing_bias_grad.dtype == torch.float32 + ), f"bias_grad must be torch.float32 (kernel uses atomic_add which requires float32), got {existing_bias_grad.dtype}." + assert existing_bias_grad.stride() == ( + K, + 1, + ), f"bias_grad must be row-major with stride (K, 1) = ({K}, 1), got {existing_bias_grad.stride()}." + + # Always zero the tensor since bias_grad represents gradients for the current + # computation and should start fresh. The kernel uses atomic_add which adds to + # existing values, so we must zero before the kernel runs. + existing_bias_grad.zero_() + + return existing_bias_grad + + else: + return gen_tgmm_bias_grad(K, G, device=device, with_bias_grad=False) + + +def get_tgmm_transposition(lhs: Tensor, rhs: Tensor, out: Tensor) -> tuple[bool, int]: + assert lhs.dim() == 2, f"lhs must have 2 dimensions (it's {lhs.dim()})." + assert rhs.dim() == 2, f"rhs must have 2 dimensions (it's {rhs.dim()})." + assert out.dim() == 3, f"out must have 3 dimensions (it's {out.dim()})." + + lhs_k, lhs_m = lhs.shape + rhs_m, rhs_n = rhs.shape + G, out_k, out_n = out.shape + + assert ( + lhs_m == rhs_m + ), f"M dimension of lhs and rhs don't match (lhs = {lhs_m}, rhs = {rhs_m})." + M = lhs_m + assert ( + lhs_k == out_k + ), f"K dimension of lhs and out don't match (lhs = {lhs_k}, rhs = {out_k})." + K = lhs_k + assert ( + rhs_n == out_n + ), f"N dimension of rhs and out don't match (lhs = {rhs_n}, rhs = {out_n})." + N = rhs_n + + assert M > 0, f"M must be positive, it's {M}." + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}" + assert G > 0, f"G must be positive, it's {G}" + + is_lhs_row_major = lhs.stride() == (M, 1) + is_lhs_col_major = lhs.stride() == (1, K) + assert ( + is_lhs_row_major != is_lhs_col_major + ), "lhs must be row-major or column-major." + is_rhs_row_major = rhs.stride() == (N, 1) + assert is_rhs_row_major, "rhs must be row-major." + is_out_row_major = out.stride() == (K * N, N, 1) + assert is_out_row_major, "out must be row-major." + + # Get lhs leading dimension according to transposition configuration. + ld_lhs = M if is_lhs_row_major else K + + return is_lhs_col_major, ld_lhs diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/logger.py b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..391ddf9b6543f5244e7f4932c8568d60748e15cd --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_grouped_gemm_triton/utils/logger.py @@ -0,0 +1,47 @@ +import os +import logging + + +# AITER Triton Logger which is singleton object around python logging. +# Note: Python logging is also a singleton object, but we want to read the +# env var AITER_LOG_LEVEL once at the beginning. Another alternative is to do +# this in __init__.py. In fact, that's how CK logger is setup. We can look at +# switching to that at some point +# +# AITER_LOG_LEVEL follows python logging levels +# DEBUG +# INFO +# WARNING +# ERROR +# CRITICAL +# +class AiterTritonLogger(object): + _instance = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super(AiterTritonLogger, cls).__new__(cls) + log_level_str = os.getenv("AITER_TRITON_LOG_LEVEL", "WARNING").upper() + numeric_level = getattr(logging, log_level_str, logging.WARNING) + cls._instance._logger = logging.getLogger("AITER_TRITON") + cls._instance._logger.setLevel(numeric_level) + + return cls._instance + + def get_logger(self): + return self._logger + + def debug(self, msg): + self._logger.debug(msg) + + def info(self, msg): + self._logger.info(msg) + + def warning(self, msg): + self._logger.warning(msg) + + def error(self, msg): + self._logger.error(msg) + + def critical(self, msg): + self._logger.critical(msg) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a720e7a2cc4e44636f6e433a2750e945dc38e8b2 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/__init__.py @@ -0,0 +1,10 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +# from megablocks.layers.dmoe import dMoE +from .moe import MoE + +__all__ = [ + 'MoE', + # 'dMoE', +] diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/activation_fn.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/activation_fn.py new file mode 100644 index 0000000000000000000000000000000000000000..0e1d956704840aa4daf7d1d71d24e051567feab9 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/activation_fn.py @@ -0,0 +1,33 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Callable, Union + +import torch +from ..stk import Matrix + + +def act_fn( + x: Matrix, + function: Callable, + return_grad_fn: bool = False, + **kwargs, +) -> Union[tuple[Matrix, Any] | Matrix]: + assert isinstance(x, Matrix) + with torch.set_grad_enabled(torch.is_grad_enabled() or return_grad_fn): + if return_grad_fn: + x.data.requires_grad = True + out = function(x.data, **kwargs) + y = Matrix( + x.size(), + out, + x.row_indices, + x.column_indices, + x.offsets, + x.column_indices_t, + x.offsets_t, + x.block_offsets_t, + ) + if return_grad_fn: + return y, out.backward + return y diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/all_to_all.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/all_to_all.py new file mode 100644 index 0000000000000000000000000000000000000000..5ac7067bcaa34db1d82b340c43550fe3577aa7a3 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/all_to_all.py @@ -0,0 +1,54 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.distributed as dist + + +class AllToAllOp(torch.autograd.Function): + + @staticmethod + def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op): + out = torch.empty((sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype) + + ctx.input_shape = x.shape + ctx.output_split_sizes = output_split_sizes + ctx.input_split_sizes = input_split_sizes + ctx.group = group + handle = dist.all_to_all_single( + out, + x, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group, + async_op=async_op, + ) + return out, handle + + @staticmethod + def backward(ctx, grad, _): + if ctx.needs_input_grad[0]: + out = torch.empty( + ctx.input_shape, + device=grad.device, + dtype=grad.dtype, + ) + dist.all_to_all_single( + out, + grad, + output_split_sizes=ctx.input_split_sizes, + input_split_sizes=ctx.output_split_sizes, + group=ctx.group, + ) + return out, None, None, None, None + return None, None, None, None, None + + +def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False): + return AllToAllOp.apply( + x, + output_split_sizes, + input_split_sizes, + group, + async_op, + ) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/arguments.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/arguments.py new file mode 100644 index 0000000000000000000000000000000000000000..4db9b1bd38bc2e2f421625c124f86b85f45c5ae0 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/arguments.py @@ -0,0 +1,101 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +from functools import partial +from typing import Any, Callable, Optional, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +# import megablocks.grouped_gemm_util as grouped_gemm +from .. import grouped_gemm_util as grouped_gemm + +# Type annotation for in-place Tensor initialization function. +InitFn = Union[Callable[[torch.Tensor], None], partial[torch.Tensor]] + +_ALLOWED_BITWIDTHS = (-1, 4, 8) + +DEFAULT_ACTIVATION_FN = partial(F.gelu, approximate='tanh') + + +@dataclasses.dataclass +class Arguments: + # Model arguments. + hidden_size: int = 1024 + ffn_hidden_size: int = 4096 + num_layers: int = 1 + bias: bool = True + return_bias: bool = True + activation_fn: Optional[Callable] = DEFAULT_ACTIVATION_FN + + # MoE arguments. + moe_num_experts: int = 1 + moe_top_k: int = 1 + moe_capacity_factor: int = 1 + moe_normalize_expert_weights: Optional[Union[int, float]] = None + moe_loss_weight: float = 0.1 + moe_jitter_eps: Optional[float] = None + moe_lbl_in_fp32: bool = False + + # Parallelism arguments. + moe_expert_model_parallelism: bool = False + expert_parallel_group: Optional[dist.ProcessGroup] = None + pipeline_model_parallel_size: int = 1 + num_layers_per_virtual_pipeline_stage: Optional[int] = None + + # Compute arguments. + memory_optimized_mlp: bool = False + mlp_type: str = 'mlp' + mlp_impl: str = 'sparse' + + # Initialization arguments. + fp16: bool = True + bf16: bool = False + device: Union[int, torch.device] = dataclasses.field(default_factory=torch.cuda.current_device) + init_method: InitFn = partial(torch.nn.init.normal_, mean=0.0, std=0.02) + output_layer_init_method: InitFn = init_method + + # Benchmarking arguments. + uniform_expert_assignment: bool = False + + # shared expert arguments + shared_expert: bool = False # enable using shared expert + fc_cls: Any = torch.nn.Linear # class of the fully connected layer in shared expert (purpose: to allow using custom FC layer eg te.Linear (for FP8)) + fc_kwargs: dict[str, Any] = dataclasses.field(default_factory=dict,) # kwargs for custom fc layers + remat_act_fn: bool = True # enable act fn to be rematerialized instead of stored + shared_expert_hidden_size: Optional[ + int] = None # hidden size of the shared expert IF we want to set it to something different from hidden_size + shared_expert_weighted_sum: bool = False # enable using weighted sum for shared expert output (wieghted by number of experts used) + + # Router Z-loss arguments + moe_zloss_weight: float = 0 # 1e-3 is a reasonable value + moe_zloss_in_fp32: bool = False + + def __post_init__(self): + # Sparse MLP is not supported with triton >=3.2.0 + # TODO: Remove this once sparse is supported with triton >=3.2.0 + if self.__getattribute__('mlp_impl') == 'sparse': + try: + import triton + if triton.__version__ >= '3.2.0': + raise ValueError( + 'Sparse MLP is not supported with triton >=3.2.0. Please use mlp_impl="grouped" instead.', + ) + except ImportError: + raise ImportError('Triton is required for sparse MLP implementation') + + if self.__getattribute__('mlp_impl') == 'grouped': + grouped_gemm.assert_grouped_gemm_is_available() + + if self.shared_expert_hidden_size is None: + self.shared_expert_hidden_size = self.ffn_hidden_size + + +def from_megatron(megatron_args: Any): + args = Arguments() + for field in dataclasses.fields(args): + if hasattr(megatron_args, field.name): + setattr(args, field.name, getattr(megatron_args, field.name)) + return args diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/common.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/common.py new file mode 100644 index 0000000000000000000000000000000000000000..2d07109702963ba48a3b94ab860807954dfd79c1 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/common.py @@ -0,0 +1,26 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from .arguments import Arguments + + +def dtype(args: Arguments): + if args.fp16: + return torch.float16 + elif args.bf16: + return torch.bfloat16 + return None + + +def cast_if_autocast_enabled(tensor): + if torch.is_autocast_enabled(): + if tensor.device.type == 'cuda': + dtype = torch.get_autocast_gpu_dtype() + elif tensor.device.type == 'cpu': + dtype = torch.get_autocast_cpu_dtype() + else: + raise NotImplementedError() + return tensor.to(dtype=dtype) + return tensor diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/dmlp_registry.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/dmlp_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..de2ed047042e438c7190ebb139b6f7f30009734c --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/dmlp_registry.py @@ -0,0 +1,42 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Union + +from . import glu, mlp +from .arguments import Arguments + +MlpType = Union[mlp.SparseMLP, glu.SparseGLU] + +_REGISTRY = { + 'mlp': { + 'grouped': mlp.GroupedMLP, + 'sparse': mlp.SparseMLP, + }, + 'glu': { + 'grouped': glu.GroupedGLU, + 'sparse': glu.SparseGLU, + }, +} + + +def get(args: Arguments) -> MlpType: + """Returns an MLP for use in a dMoE instance. + + Uses the provided arguments to instantiate the appropriate + MLP instance. This only contains MLPs for use in dMoEs + (ie. only for the dropless versions of MoEs). + + Args: + args: propagated Arguments dataclass. + + Returns: + An instantiated MLP constructed using the input args. + """ + if args.mlp_type not in _REGISTRY: + raise ValueError(f'Unsupported mlp type: {args.mlp_type}') + + if args.mlp_impl not in _REGISTRY[args.mlp_type]: + raise ValueError(f'{args.mlp_type} does not support {args.mlp_impl} backend.',) + + return _REGISTRY[args.mlp_type][args.mlp_impl](args) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/dmoe.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/dmoe.py new file mode 100644 index 0000000000000000000000000000000000000000..6d0375a4df2f27134c4127e60be04f3b45693050 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/dmoe.py @@ -0,0 +1,337 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import torch + +# try: +# import stk.ops +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/_layers/dmoe.py is needed.', +# ) + +# import megablocks.ops as ops +# # from megablocks.ops import ops +# from megablocks.layers import common, dmlp_registry, moe, mpu +# from megablocks.layers.arguments import Arguments + +from .. import stk +from .. import ops +from . import common, dmlp_registry, moe, mpu +from .arguments import Arguments + +def promote_scalar(x): + return x.view(1) if not len(x.size()) else x + + +class ParallelDroplessMLP(moe.ParallelMLP): + + def __init__(self, args: Arguments): + super(ParallelDroplessMLP, self).__init__(args) + self.hidden_size = args.hidden_size + self.ffn_hidden_size = mpu.features_per_rank(args) + self.blocking = 128 + self.mlp = dmlp_registry.get(args) + + # Calculate the number of bits needed to represent the column indices + # in the intermediate sparse matrix. + max_column_index = ((self.ffn_hidden_size * self.num_experts) // self.blocking) + self.transpose_sort_end_bit = max( + int(np.ceil(np.log2(max_column_index))), + 1, + ) + + def sparse_transpose(self, size, row_indices, column_indices, offsets): + block_columns = size[1] // self.blocking + + # Sort row indices by column indices to get the transposed matrix's + # column indices. + # + # NOTE: Our sort operation uses the same width indices as the input values. + # To avoid overflow when we have large activation matrices we cast to + # 32-bit before sorting. + _, gather_indices = ops.sort( + column_indices.int(), + self.transpose_sort_end_bit, + ) + + # There are a constant number of blocks in every row of the sparse matrix. + # A blocks offset is: + # + # row_index * blocks_per_row + column_index % blocks_per_row + # + # Once we have the block offsets ordered for transposition we can divide + # by blocks_per_row to get the transposed column indices. + column_indices_t = row_indices.gather(0, gather_indices.long()) + block_offsets_t = gather_indices.int() + + zero = torch.zeros((1,), dtype=torch.int32, device=row_indices.device) + nnz_per_column = ops.histogram(column_indices, block_columns) + nnz_per_column = ops.inclusive_cumsum(nnz_per_column, 0) + if nnz_per_column.dim() == 0: + # This addresses an edge case when ffn_hidden_size is equal to self.blocking. + nnz_per_column = nnz_per_column.unsqueeze(0) + offsets_t = torch.cat([zero, nnz_per_column]) + return column_indices_t, offsets_t, block_offsets_t + + def topology(self, x, padded_bins): + padded_tokens, _ = x.size() + assert padded_tokens % self.blocking == 0 + if self.ffn_hidden_size % self.blocking != 0: + raise ValueError( + f'The ffn_hidden_size {self.ffn_hidden_size} must be divisible by ' + + f'the block size {self.blocking}. Please update your configuration.', + ) + + # Offsets for the sparse matrix. All rows have the + # same number of nonzero blocks dictated by the + # dimensionality of a single expert. + block_rows = padded_tokens // self.blocking + blocks_per_row = self.ffn_hidden_size // self.blocking + offsets = torch.arange( + 0, + block_rows * blocks_per_row + 1, + blocks_per_row, + dtype=torch.int32, + device=x.device, + ) + + # Indices for the sparse matrix. The indices for + # the intermediate matrix are dynamic depending + # on the mapping of tokens to experts. + column_indices = ops.topology( + padded_bins, + self.blocking, + block_rows, + blocks_per_row, + ) + + # TODO(tgale): This is unused. Remove the need for this in stk. + # For now, use meta init to save the device memory. + data = torch.empty( + column_indices.numel(), + self.blocking, + self.blocking, + dtype=common.dtype(self.args), + device='meta', + ) + shape = ( + padded_tokens, + self.ffn_hidden_size * mpu.experts_per_rank(self.args), + ) + row_indices = stk.ops.row_indices(shape, data, offsets, column_indices) + column_indices_t, offsets_t, block_offsets_t = self.sparse_transpose( + shape, + row_indices, + column_indices, + offsets, + ) + return stk.Matrix( + shape, + data, + row_indices, + column_indices, + offsets, + column_indices_t, + offsets_t, + block_offsets_t, + ) + + def indices_and_padded_bins(self, top_experts): + # Sort the expert ids to produce the scatter/gather + # indices for the permutation. + top_experts = top_experts.int() + bin_ids, indices = ops.sort(top_experts, self.sort_end_bit) + + # Histogram the expert ids to identify the number of + # tokens routed to each expert. + tokens_per_expert = ops.histogram(top_experts, self.num_experts) + + # Round the token counts up to the block size used in + # the matrix muliplications. Caculate the starting + # position of each bin. + padded_tokens_per_expert = ops.round_up( + tokens_per_expert, + self.blocking, + ) + padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) + padded_bins = promote_scalar(padded_bins) + + # Calculate the bin bounds for the sorted tokens. + bins = ops.inclusive_cumsum(tokens_per_expert, 0) + bins = promote_scalar(bins) + return indices, bin_ids, bins, padded_bins, tokens_per_expert + + def sparse_forward_once(self, x, expert_weights, top_experts): + # x: [sl, bs, hs] + # expert_weights: [sl * bs, top-k] + # top_experts: [sl * bs, top-k] + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + with torch.no_grad(): + indices, bin_ids, bins, padded_bins, tokens_per_expert = (self.indices_and_padded_bins(top_experts)) + + # Route the tokens for MoE computation. + x = x.view(-1, x.shape[-1]) + x = ops.padded_gather( + x, + indices, + bin_ids, + bins, + padded_bins, + self.top_k, + ) + + # Create the sparse matrix topology. + with torch.no_grad(): + topo = self.topology(x, padded_bins) + + # Perform the expert computation. + x = self.mlp(x, topo) + + # Un-route the data for the MoE output. + x = ops.padded_scatter( + x, + indices, + bin_ids, + expert_weights, + bins, + padded_bins, + self.top_k, + ) + return x, tokens_per_expert + + # For use in the base-class parallel_forward_once. + def sparse_permute_and_compute( + self, + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, # unused + top_k, + ): + + # Round the token counts up to the block size used in the matrix + # multiplication. Calculate the starting position of each bin. + padded_tokens_per_expert = ops.round_up( + tokens_per_expert, + self.blocking, + ) + padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) + padded_bins = promote_scalar(padded_bins) + + # Route the tokens for MoE computation. + x = x.view(-1, x.shape[-1]) + x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k) + + # Create the sparse matrix topology. + with torch.no_grad(): + topo = self.topology(x, padded_bins) + + # Perform the expert computation. + x = self.mlp(x, topo) + + # Un-route the data for the MoE output. + return ops.padded_scatter( + x, + indices, + bin_ids, + expert_weights, + bins, + padded_bins, + top_k, + ) + + def grouped_forward_once(self, x, expert_weights, top_experts): + # x: [sl, bs, hs] + # expert_weights: [sl * bs, top-k] + # top_experts: [sl * bs, top-k] + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + with torch.no_grad(): + indices, bin_ids, bins, tokens_per_expert = (self.indices_and_bins(top_experts)) + + out = self.grouped_permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + -1, # unused + self.args.moe_top_k, + ) + return out, tokens_per_expert + + def grouped_permute_and_compute( + self, + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, # unused + top_k, + ): + + # Route the tokens for MoE computation. + x = x.view(-1, x.shape[-1]) + x = ops.gather(x, indices, bin_ids, bins, top_k) + + # Perform the expert computation. + x = self.mlp(x, tokens_per_expert) + + # Un-route the data for the MoE output. + return ops.scatter(x, indices, bin_ids, expert_weights, bins, top_k) + + def forward_once(self, x, expert_weights, top_experts): + if self.args.mlp_impl == 'sparse': + return self.sparse_forward_once(x, expert_weights, top_experts) + else: + return self.grouped_forward_once(x, expert_weights, top_experts) + + def permute_and_compute( + self, + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, + top_k, + ): + if self.args.mlp_impl == 'sparse': + return self.sparse_permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, + top_k, + ) + else: + return self.grouped_permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, + top_k, + ) + + +class dMoE(moe.MoE): + + def _init_experts_mlp(self, args: Arguments): + return ParallelDroplessMLP(args) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/gelu.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/gelu.py new file mode 100644 index 0000000000000000000000000000000000000000..c4c9e6532798615b5c12c96694241a4c18ee8f7b --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/gelu.py @@ -0,0 +1,52 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +# try: +# import stk +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/_layers/gelu.py is needed.', +# ) + +from .. import stk + +import torch +import torch.nn.functional as F + + +@torch.jit.script +def _gelu_backward_inplace(g, x): + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + ff = (0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)) + return g.mul_(ff) + + +def gelu_backward_(grad: stk.Matrix, x: stk.Matrix): + # NOTE: The two sparse matrices must have the same topology. + if isinstance(grad, stk.Matrix) and isinstance(x, stk.Matrix): + return stk.Matrix( + x.size(), + _gelu_backward_inplace(grad.data, x.data), + x.row_indices, + x.column_indices, + x.offsets, + x.column_indices_t, + x.offsets_t, + x.block_offsets_t, + ) + return _gelu_backward_inplace(grad, x) + + +def gelu(x: stk.Matrix): + assert isinstance(x, stk.Matrix) + return stk.Matrix( + x.size(), + F.gelu(x.data, approximate='tanh'), + x.row_indices, + x.column_indices, + x.offsets, + x.column_indices_t, + x.offsets_t, + x.block_offsets_t, + ) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/glu.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/glu.py new file mode 100644 index 0000000000000000000000000000000000000000..5f297a41ff6a1a2a285f5b461951672364b898da --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/glu.py @@ -0,0 +1,244 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +# import stk.ops +# try: +# import stk.ops +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/_layers/glu.py is needed.', +# ) + +from .. import stk + +import torch + +# from megablocks import grouped_gemm_util as gg +# from megablocks.layers import common, mpu +# from megablocks.layers.activation_fn import act_fn +# from megablocks.layers.arguments import Arguments +# from megablocks.layers.mlp import ( +# SharedMLP, +# SparseMLP, +# create_dmoe_expert_weights, +# resolve_dtensor, +# ) + +from .. import grouped_gemm_util as gg +from . import common, mpu +from .activation_fn import act_fn +from .arguments import Arguments +from .mlp import ( + SharedMLP, + SparseMLP, + create_dmoe_expert_weights, + resolve_dtensor, +) + + +class SparseGLU(SparseMLP): + + def __init__(self, args: Arguments): + super().__init__(args) + self.v1 = torch.nn.Parameter( + torch.empty( + self._num_rows_per_rank, + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + with torch.no_grad(): + self.v1.copy_( + create_dmoe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.init_method, + ), + ) + + mpu.set_expert_model_parallel_attributes( + self.v1, + self._should_set_parallelism_attribute, + ) + + def forward(self, x, topo): + if self.args.memory_optimized_mlp: + raise NotImplementedError( + 'Memory optimized implementation not yet supported with GLU with sparse kernels.', + ) + + w1, v1, w2 = self.scale_grad(self.w1), self.scale_grad(self.v1,), self.scale_grad(self.w2) + w1, v1, w2 = resolve_dtensor(w1), resolve_dtensor(v1,), resolve_dtensor(w2) + + # Compute the GLU. + x1 = stk.ops.sdd(x, w1.t(), topo) + x2 = stk.ops.sdd(x, v1.t(), topo) + + activation_fn_out = act_fn(x1, self.args.activation_fn) + x1 = stk.ops.mul(activation_fn_out, x2) + + return stk.ops.dsd(x1, w2) + + +class MemoryOptimizedGroupedGLU(torch.autograd.Function): + """GroupedMLP with manually scheduled memory reuse.""" + + @staticmethod + @torch.amp.autocast_mode.custom_fwd(device_type='cuda') + def forward(ctx, x, w1, v1, w2, batch_sizes, activation_fn): + # Cast inputs using ctx dtype from AMP + if ctx._fwd_used_autocast: + x = x.to(ctx._dtype) + w1 = w1.to(ctx._dtype) + v1 = v1.to(ctx._dtype) + w2 = w2.to(ctx._dtype) + # x: [m, k], w1: [n, k], v1: [n, k], w2: [n, k] + if (not x.is_contiguous() or not w1.is_contiguous() or not v1.is_contiguous() or not w2.is_contiguous()): + raise ValueError("Expected contiguous 'x', 'w1', 'v1' and 'w2'.") + + # Layer 0: x @ w1.t(). + assert gg.backend is not None + sdd_out = gg.backend.gmm(x, w1, batch_sizes, trans_b=True) + v1_out = gg.backend.gmm(x, v1, batch_sizes, trans_b=True) + + # GeLU. + activation_fn_out = activation_fn(sdd_out) * v1_out + + # Layer 1: x @ w2. + dsd_out = gg.backend.gmm(activation_fn_out, w2, batch_sizes) + + # NOTE: Save the input to the layer and the activation_fn input for + # gradient computation. We'll re-compute the activation_fn forward + # pass in the backward pass to avoid materializing another + # intermediate. + ctx.x_shape = x.shape + ctx.sdd_out_shape = sdd_out.shape + ctx.dtype = x.dtype + ctx.activation_fn = activation_fn + ctx.save_for_backward(w1, v1, w2, batch_sizes, x, sdd_out, v1_out) + return dsd_out + + @staticmethod + @torch.amp.autocast_mode.custom_bwd(device_type='cuda') + def backward(ctx, ddsd_out): + if (not ctx.needs_input_grad[0] or not ctx.needs_input_grad[1] or not ctx.needs_input_grad[2]): + raise ValueError('Expected all MLP inputs to need grad.') + + # Unpack saved tensors + # dtype = ctx.dtype + saved_tensors = ctx.saved_tensors + w1, v1, w2 = saved_tensors[:3] + batch_sizes = saved_tensors[3] + x = saved_tensors[4] + sdd_out, v1_out = saved_tensors[5:7] + + # Rematerialize activation_fn output. + activation_fn = ctx.activation_fn + with torch.set_grad_enabled(True): + sdd_out.requires_grad = True + v1_out.requires_grad = True + activation_fn_out = activation_fn(sdd_out) * v1_out + activation_grad_fn = activation_fn_out.backward + + # Compute dw2 with recomputed activation_fn output. + assert gg.backend is not None + dw2 = gg.backend.gmm( + activation_fn_out, + ddsd_out, + batch_sizes, + trans_a=True, + ) + + # Compute dactivation_fn_out. + # + # NOTE: We reuse the activation_fn_out allocation. + dactivation_fn_out = activation_fn_out + gg.backend.gmm( + ddsd_out, + w2, + batch_sizes, + trans_b=True, + c=dactivation_fn_out, + ) + + # Compute dsdd_out. + # + # NOTE: This reuses the dactivation_fn_out allocation. + assert activation_grad_fn is not None + activation_grad_fn(dactivation_fn_out) + dsdd_out = sdd_out.grad + dv1_out = v1_out.grad + + # Compute dw1. + dw1 = gg.backend.gmm(dsdd_out, x, batch_sizes, trans_a=True) + + # Compute dv1. + dv1 = gg.backend.gmm(dv1_out, x, batch_sizes, trans_a=True) + + # Compute dx. + # + # NOTE: This reuses the ddsd_out allocation. + dx = ddsd_out + gg.backend.gmm(dsdd_out, w1, batch_sizes, c=dx) + dx += gg.backend.gmm(dv1_out, v1, batch_sizes) + return dx, dw1, dv1, dw2, None, None + + +memory_optimized_grouped_glu = MemoryOptimizedGroupedGLU.apply + + +class GroupedGLU(SparseGLU): + + def forward(self, x, tokens_per_expert): + batch_sizes = tokens_per_expert.cpu().to(torch.long) + w1, v1, w2 = ( + self.scale_grad(self.w1), + self.scale_grad(self.v1), + self.scale_grad(self.w2), + ) + w1, v1, w2 = resolve_dtensor(w1), resolve_dtensor(v1,), resolve_dtensor(w2) + + # Re-shape the weights for the grouped GEMMs. + ne = mpu.experts_per_rank(self.args) + w1 = w1.view(ne, -1, self.args.hidden_size) + v1 = v1.view(ne, -1, self.args.hidden_size) + w2 = w2.view(ne, -1, self.args.hidden_size) + + if self.args.memory_optimized_mlp: + return memory_optimized_grouped_glu( + x, + w1, + v1, + w2, + batch_sizes, + self.args.activation_fn, + ) + + # Compute the MLP. + assert gg.ops is not None + x1 = gg.ops.gmm(x, w1, batch_sizes, trans_b=True) + x2 = gg.ops.gmm(x, v1, batch_sizes, trans_b=True) + x1 = self.args.activation_fn(x1) * x2 + return gg.ops.gmm(x1, w2, batch_sizes) + + +class SharedGLU(SharedMLP): + """GPU for shared expert. + + Note: this is a copy -> pasta -> modify of the LLM-Foundry MPTGLU class + """ + + def __init__(self, args: Arguments): + super().__init__(args) + self.gate_proj = args.fc_cls( + args.hidden_size, + self.args.shared_expert_hidden_size, + **self.fc_kwargs, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x)) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/memory_test.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/memory_test.py new file mode 100644 index 0000000000000000000000000000000000000000..74d1166931b712635131985b25a89f4ca23e576d --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/memory_test.py @@ -0,0 +1,103 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import gc + +import torch +import torch.distributed as dist + +# from megablocks.layers import arguments, dmoe +from . import arguments, dmoe + +_TESTS = ((8, 2048, 4096, 4096, 32, 4),) + + +def get_tensors(): + ptrs = set() + out = [] + for obj in gc.get_objects(): + if torch.is_tensor(obj): + if not obj.is_contiguous() or obj.data_ptr() in ptrs: + continue + out.append(obj) + ptrs.add(obj.data_ptr()) + return out + + +def test_memory( + group, + batch_size, + sequence_length, + hidden_size, + ffn_hidden_size, + num_experts, + top_k, +): + args = arguments.Arguments( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + moe_num_experts=num_experts, + moe_top_k=top_k, + moe_expert_model_parallelism=True, + expert_parallel_group=group, + fp16=False, + bf16=True, + device=torch.cuda.current_device(), + ) + layer = dmoe.dMoE(args).cuda() + + x = torch.randn((batch_size, sequence_length, hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16).requires_grad_(True) + torch.cuda.empty_cache() + + # Run forward + backward. + # with torch.autograd.detect_anomaly(): + out, _ = layer(x) + out.mean().backward() + + # Report peak memory. + mem = torch.cuda.max_memory_allocated() + print('Max Memory Allocated = {:0.0f}MiB'.format(mem / 1e6)) + print('Max Memory Reserved = {:0.0f}MiB'.format(torch.cuda.max_memory_reserved() / 1e6,),) + + # Calculate weight and gradient memory usage. + weight_memory = 2 * ( + layer.router.layer.weight.numel() + layer.experts.mlp.w1.numel() + layer.experts.mlp.w2.numel() + ) + + def grad_numel(x): + if x.grad is not None: + return x.grad.numel() + return 0 + + grad_memory = 2 * ( + grad_numel(layer.router.layer.weight) + grad_numel(layer.experts.mlp.w1) + grad_numel(layer.experts.mlp.w2) + ) + weight_memory += grad_memory + + print('Weight Memory Allocated = {:0.0f}MiB'.format(weight_memory / 1e6)) + print('Activation Memory Allocated = {:0.0f}MiB'.format((mem - weight_memory) / 1e6,),) + + # Manually calculate GPU memory usage from the garbage + # collector. + gc.collect() + total = 0 + tensors = get_tensors() + tensors = sorted(tensors, key=lambda x: -x.numel()) + for i, t in enumerate(tensors): + total += t.numel() + print(f'{i}: {t.shape}, {t.numel() * 2}') + del tensors + + print('Total Bytes Found = {:0.0f}MiB'.format(total * 2 / 1e6)) + + +if __name__ == '__main__': + assert dist.is_available() + group = dist.init_process_group(backend='nccl') + local_rank = dist.get_rank(group) + torch.cuda.set_device(local_rank) + + for args in _TESTS: + test_memory(group, *args) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/mlp.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..c99afb9904c24a8b6a83e79059cd1251dbbfd99e --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/mlp.py @@ -0,0 +1,587 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# try: +# import stk +# import stk.backend.triton_kernels +# import stk.ops +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/_layers/mlp.py is needed.', +# ) + +from .. import stk + +import torch +from packaging import version + +# from megablocks import grouped_gemm_util as gg +# from megablocks.layers import common, gelu, mpu +# from megablocks.layers.activation_fn import act_fn +# from megablocks.layers.arguments import DEFAULT_ACTIVATION_FN, Arguments, InitFn + +from .. import grouped_gemm_util as gg +from . import common, gelu, mpu +from .activation_fn import act_fn +from .arguments import DEFAULT_ACTIVATION_FN, Arguments, InitFn + +class ScaleGradient(torch.autograd.Function): + + @staticmethod + @torch.amp.autocast_mode.custom_fwd(device_type='cuda') + def forward(ctx: Any, x: torch.Tensor, scale: float): + ctx.scale = scale + return x + + @staticmethod + @torch.amp.autocast_mode.custom_bwd(device_type='cuda') + def backward(ctx: torch.Tensor, grad: torch.Tensor): + return grad * ctx.scale, None + + +scale_gradient = ScaleGradient.apply + + +def resolve_dtensor(weight: torch.Tensor): + if version.parse(torch.__version__) >= version.parse('2.0.0'): + from torch.distributed._tensor import DTensor + if isinstance(weight, DTensor): + return weight.to_local() + return weight + + +def create_moe_expert_weights( + args: Arguments, + num_experts: int, + ffn_hidden_size: int, + hidden_size: int, + init_method: InitFn, +): + # Create the entire weight matrix such that the sampled weights will + # not vary between data parallelism and expert model parallelism for + # the same random seed. + master_weights = torch.empty( + num_experts, + ffn_hidden_size, + hidden_size, + device=args.device, + dtype=common.dtype(args), + ) + init_method(master_weights) + + if not args.moe_expert_model_parallelism: + return master_weights + + # Calculate the amount of sharding in each dimension. + expert_sharding_degree = mpu.expert_sharding_degree(args) + hidden_sharding_degree = mpu.hidden_sharding_degree(args) + + # Calculate the experts per rank. + # + # NOTE: We assign ranks to be expert parallel before going + # tensor parallel. + rank = mpu.get_expert_parallel_rank(args) + expert_rank = rank % expert_sharding_degree + num_experts_per_rank = num_experts // expert_sharding_degree + start_expert = expert_rank * num_experts_per_rank + end_expert = (expert_rank + 1) * num_experts_per_rank + + # Calculate the rows per rank. + row_rank = rank // expert_sharding_degree + num_rows_per_rank = ffn_hidden_size // hidden_sharding_degree + start_row = row_rank * num_rows_per_rank + end_row = (row_rank + 1) * num_rows_per_rank + + # Slice the weight matrix to get the chunk for this rank. + with torch.no_grad(): + weights = master_weights[start_expert:end_expert, start_row:end_row] + return weights + + +class MLP(torch.nn.Module): + + def __init__(self, args: Arguments): + super().__init__() + self.args = args + # expert_parallel_world_size = mpu.get_expert_parallel_world_size(args) + experts_per_rank = mpu.experts_per_rank(args) + + self.w1 = torch.nn.Parameter( + torch.empty( + experts_per_rank, + args.hidden_size, + mpu.features_per_rank(args), + device=args.device, + dtype=common.dtype(args), + ), + ) + self.w2 = torch.nn.Parameter( + torch.empty( + experts_per_rank, + mpu.features_per_rank(args), + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + mpu.set_expert_model_parallel_attributes( + self.w1, + args.moe_expert_model_parallelism, + ) + mpu.set_expert_model_parallel_attributes( + self.w2, + args.moe_expert_model_parallelism, + ) + + # Initialize the parameters for the MLP. + # + # NOTE: It is important that we create the weight tensors prior + # to creating the master weights and slicing our the piece for + # this rank. If the master weights are created first the PyTorch + # caching allocator appears to use the same memory block for these + # and the slice which causes large increases in our peak memory + # usage. + with torch.no_grad(): + w1 = create_moe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.init_method, + ) + self.w1.copy_(w1.transpose(1, 2).contiguous()) + self.w2.copy_( + create_moe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.output_layer_init_method, + ), + ) + + self.gradient_scale = None + if self.args.moe_expert_model_parallelism: + self.gradient_scale = 1 / mpu.get_expert_parallel_world_size(self.args,) + + def scale_grad(self, w): + if self.gradient_scale is None: + return w + return scale_gradient(w, self.gradient_scale) + + def forward(self, x): + w1, w2 = self.scale_grad(self.w1), self.scale_grad(self.w2) + w1, w2 = resolve_dtensor(w1), resolve_dtensor(w2) + x = torch.bmm(x, w1) + x = self.args.activation_fn(x) + return torch.bmm(x, w2) + + +def create_dmoe_expert_weights( + args: Arguments, + num_experts: int, + rows: int, + columns: int, + init_method: InitFn, +): + weights = create_moe_expert_weights( + args, + num_experts, + rows, + columns, + init_method, + ) + return weights.view([-1, columns]) + + +class MemoryOptimizedMLP(torch.autograd.Function): + """Sparse MLP with manually scheduled memory reuse.""" + + @staticmethod + @torch.amp.autocast_mode.custom_fwd(device_type='cuda') + def forward(ctx, x, w1, w2, topo, activation_fn): + # Cast inputs using ctx dtype from AMP + if ctx._fwd_used_autocast: + x = x.to(ctx._dtype) + w1 = w1.to(ctx._dtype) + w2 = w2.to(ctx._dtype) + # x: [m, k], w1: [n, k], w2: [n, k] + if (not x.is_contiguous() or not w1.is_contiguous() or not w2.is_contiguous()): + raise ValueError("Expected contiguous 'x', 'w1' and 'w2'.") + + topo_tensors = ( + topo.row_indices, + topo.column_indices, + topo.offsets, + topo.column_indices_t, + topo.offsets_t, + topo.block_offsets_t, + ) + + # Layer 0: x @ w1.t(). + sdd_out = stk.ops.sdd(x, w1.t(), topo) + + # GeLU. + activation_fn_out = act_fn(sdd_out, activation_fn) + + # Layer 1: x @ w2. + dsd_out = stk.ops.dsd(activation_fn_out, w2) + + # NOTE: Save the input to the layer and the activation_fn input for + # gradient computation. We'll re-compute the activation_fn forward + # pass in the backward pass to avoid materializing another + # intermediate. + ctx.shape = topo.shape + ctx.x_shape = x.shape + ctx.sdd_out_shape = sdd_out.data.shape + ctx.dtype = x.dtype + ctx.activation_fn = activation_fn + ctx.save_for_backward(w1, w2, *topo_tensors, x, sdd_out.data) + return dsd_out + + @staticmethod + @torch.amp.autocast_mode.custom_bwd(device_type='cuda') + def backward(ctx, ddsd_out): + if (not ctx.needs_input_grad[0] or not ctx.needs_input_grad[1] or not ctx.needs_input_grad[2]): + raise ValueError('Expected all MLP inputs to need grad.') + + # unpack saved tensors + # dtype = ctx.dtype + saved_tensors = ctx.saved_tensors + w1, w2 = saved_tensors[:2] + topo_tensors = saved_tensors[2:8] + x = saved_tensors[8] + sdd_out_data = saved_tensors[9] + + # rematerialize activation function output + activation_fn = ctx.activation_fn + sdd_out = stk.Matrix(ctx.shape, sdd_out_data, *topo_tensors) + activation_fn_out, activation_grad_fn = act_fn( + sdd_out, + activation_fn, + return_grad_fn=True, + ) + + # Compute dw2 with recomputed activation_fn output. + dw2 = stk.ops.dsd(activation_fn_out.t(), ddsd_out) + + # Compute dactivation_fn_out. + # + # NOTE: We reuse the activation_fn_out allocation. + dactivation_fn_out = activation_fn_out + stk.backend.triton_kernels.sdd( + ddsd_out, + w2.t(), + dactivation_fn_out.shape, + dactivation_fn_out.data, + dactivation_fn_out.offsets, + dactivation_fn_out.row_indices, + dactivation_fn_out.column_indices, + ) + + # Compute dsdd_out. + # + # NOTE: This reuses the dactivation_fn_out allocation. + if activation_fn is DEFAULT_ACTIVATION_FN: + dsdd_out = gelu.gelu_backward_(dactivation_fn_out, sdd_out) + else: + assert activation_grad_fn is not None + activation_grad_fn(dactivation_fn_out.data) + dsdd_out = stk.Matrix(ctx.shape, sdd_out.data.grad, *topo_tensors) + + # Compute dw1. + dw1 = stk.ops.dsd(dsdd_out.t(), x) + + # Compute dx. + # + # NOTE: This reuses the ddsd_out allocation. + stk.backend.triton_kernels.dsd( + dsdd_out.shape, + dsdd_out.data, + dsdd_out.offsets, + dsdd_out.row_indices, + dsdd_out.column_indices, + dsdd_out.offsets_t, + dsdd_out.column_indices_t, + dsdd_out.block_offsets_t, + False, + w1, + ddsd_out, + ) + dx = ddsd_out + return dx, dw1, dw2, None, None + + +memory_optimized_mlp = MemoryOptimizedMLP.apply + + +class SparseMLP(torch.nn.Module): + + def __init__(self, args: Arguments): + super().__init__() + self.args = args + self._num_rows_per_rank = mpu.experts_per_rank(args) * mpu.features_per_rank(args) + + self.w1 = torch.nn.Parameter( + torch.empty( + self._num_rows_per_rank, + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + self.w2 = torch.nn.Parameter( + torch.empty( + self._num_rows_per_rank, + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + + # Initialize the parameters for the MLP. + # + # NOTE: It is important that we create the weight tensors prior + # to creating the master weights and slicing our the piece for + # this rank. If the master weights are created first the PyTorch + # caching allocator appears to use the same memory block for these + # and the slice which causes large increases in our peak memory + # usage. + with torch.no_grad(): + self.w1.copy_( + create_dmoe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.init_method, + ), + ) + self.w2.copy_( + create_dmoe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.output_layer_init_method, + ), + ) + + self._should_set_parallelism_attribute = args.moe_expert_model_parallelism + mpu.set_expert_model_parallel_attributes( + self.w1, + self._should_set_parallelism_attribute, + ) + mpu.set_expert_model_parallel_attributes( + self.w2, + self._should_set_parallelism_attribute, + ) + + self.gradient_scale = None + if self.args.moe_expert_model_parallelism: + self.gradient_scale = 1 / mpu.get_expert_parallel_world_size(self.args,) + + def scale_grad(self, w): + if self.gradient_scale is None: + return w + return scale_gradient(w, self.gradient_scale) + + def forward(self, x, topo): + w1, w2 = self.scale_grad(self.w1), self.scale_grad(self.w2) + w1, w2 = resolve_dtensor(w1), resolve_dtensor(w2) + if self.args.memory_optimized_mlp: + return memory_optimized_mlp( + x, + w1, + w2, + topo, + self.args.activation_fn, + ) + + # Compute the MLP. + x = stk.ops.sdd(x, w1.t(), topo) + activation_fn_out = act_fn(x, self.args.activation_fn) + return stk.ops.dsd(activation_fn_out, w2) + + +class MemoryOptimizedGroupedMLP(torch.autograd.Function): + """GroupedMLP with manually scheduled memory reuse.""" + + @staticmethod + @torch.amp.autocast_mode.custom_fwd(device_type='cuda') + def forward(ctx, x, w1, w2, batch_sizes, activation_fn): + # Cast inputs using ctx dtype from AMP + if ctx._fwd_used_autocast: + x = x.to(ctx._dtype) + w1 = w1.to(ctx._dtype) + w2 = w2.to(ctx._dtype) + # x: [m, k], w1: [n, k], w2: [n, k] + if (not x.is_contiguous() or not w1.is_contiguous() or not w2.is_contiguous()): + raise ValueError("Expected contiguous 'x', 'w1' and 'w2'.") + + # Layer 0: x @ w1.t(). + assert gg.backend is not None + sdd_out = gg.backend.gmm(x, w1, batch_sizes, trans_b=True) + + # activation_fn + activation_fn_out = activation_fn(sdd_out) + + # Layer 1: x @ w2. + dsd_out = gg.backend.gmm(activation_fn_out, w2, batch_sizes) + + # NOTE: Save the input to the layer and the activation_fn input for + # gradient computation. We'll re-compute the activation_fn forward + # pass in the backward pass to avoid materializing another + # intermediate. + ctx.x_shape = x.shape + ctx.sdd_out_shape = sdd_out.shape + ctx.dtype = x.dtype + ctx.activation_fn = activation_fn + ctx.save_for_backward(w1, w2, batch_sizes, x, sdd_out) + return dsd_out + + @staticmethod + @torch.amp.autocast_mode.custom_bwd(device_type='cuda') + def backward(ctx: Any, ddsd_out: torch.Tensor): + if (not ctx.needs_input_grad[0] or not ctx.needs_input_grad[1] or not ctx.needs_input_grad[2]): + raise ValueError('Expected all MLP inputs to need grad.') + + # Unpack saved tensors + # dtype = ctx.dtype + saved_tensors = ctx.saved_tensors + w1, w2 = saved_tensors[:2] + batch_sizes = saved_tensors[2] + x = saved_tensors[3] + sdd_out = saved_tensors[4] + + # Rematerialize activation_fn output. + activation_fn = ctx.activation_fn + with torch.set_grad_enabled(True): + sdd_out.requires_grad = True + activation_fn_out = activation_fn(sdd_out) + activation_grad_fn = activation_fn_out.backward + + # Compute dw2 with recomputed activation_fn output. + assert gg.backend is not None + dw2 = gg.backend.gmm( + activation_fn_out, + ddsd_out, + batch_sizes, + trans_a=True, + ) + + # Compute dactivation_fn_out. + # + # NOTE: We reuse the activation_fn_out allocation. + dactivation_fn_out = activation_fn_out + gg.backend.gmm( + ddsd_out, + w2, + batch_sizes, + trans_b=True, + c=dactivation_fn_out, + ) + + # Compute dsdd_out. + # + # NOTE: This reuses the dactivation_fn_out allocation. + if activation_fn is DEFAULT_ACTIVATION_FN: + dsdd_out = gelu.gelu_backward_(dactivation_fn_out, sdd_out) + else: + assert activation_grad_fn is not None + activation_grad_fn(dactivation_fn_out) + dsdd_out = sdd_out.grad + + # Compute dw1. + dw1 = gg.backend.gmm(dsdd_out, x, batch_sizes, trans_a=True) + + # Compute dx. + # + # NOTE: This reuses the ddsd_out allocation. + gg.backend.gmm(dsdd_out, w1, batch_sizes, c=ddsd_out) + dx = ddsd_out + return dx, dw1, dw2, None, None + + +memory_optimized_grouped_mlp = MemoryOptimizedGroupedMLP.apply + + +class GroupedMLP(SparseMLP): + + def forward(self, x, tokens_per_expert): + batch_sizes = tokens_per_expert.cpu().to(torch.long) + w1, w2 = (self.scale_grad(self.w1), self.scale_grad(self.w2)) + + # Re-shape the weights for the grouped GEMMs. + ne = mpu.experts_per_rank(self.args) + w1 = resolve_dtensor(w1).view(ne, -1, self.args.hidden_size) + w2 = resolve_dtensor(w2).view(ne, -1, self.args.hidden_size) + + if self.args.memory_optimized_mlp: + return memory_optimized_grouped_mlp( + x, + w1, + w2, + batch_sizes, + self.args.activation_fn, + ) + + # Compute the MLP. + assert gg.ops is not None + x = gg.ops.gmm(x, w1, batch_sizes, trans_b=True) + x = self.args.activation_fn(x) + return gg.ops.gmm(x, w2, batch_sizes) + + +class SharedMLP(torch.nn.Module): + """MLP for shared expert. + + Note: this is a copy -> pasta -> modify of the LLM-Foundry MPTMLP class + """ + + def __init__(self, args: Arguments): + super().__init__() + self.args = args + self.fc_kwargs: dict[str, Any] = { + 'bias': args.bias, + 'device': args.device, + } + self.fc_kwargs.update(args.fc_kwargs) + + self.up_proj = args.fc_cls( + args.hidden_size, + args.shared_expert_hidden_size, + **self.fc_kwargs, + ) + self.act = args.activation_fn + self.down_proj = args.fc_cls( + args.shared_expert_hidden_size, + args.hidden_size, + **self.fc_kwargs, + ) + self.down_proj._is_residual = True # a flag for llm-foundry init + + def add_experts_sharedexpert( + self, + shared_expert_out: torch.Tensor, + expert_out: torch.Tensor, + ) -> torch.Tensor: + # Helper function to add expert output to shared expert output + # with optional weighted sum. + if self.args.shared_expert_weighted_sum: + # enable using weighted sum for shared expert output + # wieghted by number of experts used + t_experts = self.args.moe_top_k + 1 + sh_mlp_out = shared_expert_out / t_experts + return sh_mlp_out.add( + expert_out, + alpha=(self.args.moe_top_k / t_experts), + ) + + return shared_expert_out + expert_out + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act(self.up_proj(x))) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/moe.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/moe.py new file mode 100644 index 0000000000000000000000000000000000000000..d0a4aeaacc9c86fc70944e730c53f7a55644e05e --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/moe.py @@ -0,0 +1,507 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Optional, Tuple + +import numpy as np +import torch +import torch.distributed as dist + +# import megablocks.ops as ops +# from megablocks.layers import common, mlp, mpu, router, sharedexpert_registry +# from megablocks.layers.all_to_all import all_to_all +# from megablocks.layers.arguments import Arguments + +from ..ops import ( + sort, + histogram, + inclusive_cumsum, + exclusive_cumsum, + binned_gather, + binned_scatter, + gather, + scatter, + repeat, + replicate, +) + +from . import common, mlp, mpu, router, sharedexpert_registry +from .arguments import Arguments +from .all_to_all import all_to_all + +_LOAD_BALANCING_LOSS = [] + + +def save_load_balancing_loss(loss): + global _LOAD_BALANCING_LOSS + _LOAD_BALANCING_LOSS.append(loss) + + +def get_load_balancing_loss(): + global _LOAD_BALANCING_LOSS + return _LOAD_BALANCING_LOSS + + +def clear_load_balancing_loss(): + global _LOAD_BALANCING_LOSS + _LOAD_BALANCING_LOSS.clear() + + +def batched_load_balancing_loss(args: Arguments): + if args.moe_loss_weight == 0: + return 0.0 + + # tokens_per_expert[i].shape = (num_experts) + # expert_scores[i].shape = (tokens, num_experts) + tokens_per_expert, expert_scores = zip(*get_load_balancing_loss()) + num_layers_per_pipeline_stage = (args.num_layers // args.pipeline_model_parallel_size) + if args.num_layers_per_virtual_pipeline_stage is not None: + num_layers_per_pipeline_stage = args.num_layers_per_virtual_pipeline_stage + + if len(tokens_per_expert) != num_layers_per_pipeline_stage: + raise ValueError( + f'Expected {num_layers_per_pipeline_stage} token_per_experts ' + f'but found {len(tokens_per_expert)}.\nnum_layers = ' + f'{args.num_layers}\npipeline_model_parallel_size = ' + f'{args.pipeline_model_parallel_size}\n' + 'num_layers_per_virtual_pipeline_stage' + f' = {args.num_layers_per_virtual_pipeline_stage}', + ) + if len(expert_scores) != num_layers_per_pipeline_stage: + raise ValueError( + f'Expected {num_layers_per_pipeline_stage} expert_scores ' + f'but found {len(tokens_per_expert)}.\nnum_layers = ' + f'{args.num_layers}\npipeline_model_parallel_size = ' + f'{args.pipeline_model_parallel_size}\n' + 'num_layers_per_virtual_pipeline_stage' + f' = {args.num_layers_per_virtual_pipeline_stage}', + ) + + # Verify the shape of the tokens_per_expert and expert_scores tensors. + assert all((x.ndim == 1 and x.numel() == args.moe_num_experts for x in tokens_per_expert)) + + tokens = expert_scores[0].shape[0] + assert all(((x.ndim == 2 and x.shape[1] == args.moe_num_experts and x.shape[0] == tokens) for x in expert_scores)) + + # Concatenate the contributions of each layer and convert to + # the correct types and formats for the dot product. + expert_scores = torch.cat(expert_scores, dim=1) + if args.moe_lbl_in_fp32: + expert_scores = expert_scores.float() + if tokens != 0: + expert_scores = expert_scores.mean(dim=0) + else: + expert_scores = expert_scores.sum(dim=0) + tokens_per_expert = torch.cat(tokens_per_expert).to(expert_scores.dtype) + + expected_values = num_layers_per_pipeline_stage * args.moe_num_experts + assert tokens_per_expert.numel() == expected_values + assert expert_scores.numel() == expected_values + + # Calculate the total scale across all factors. + # + # loss_weight * num_experts / (num_layers * tokens * top_k) + scale_numerator = (args.moe_num_experts * args.moe_loss_weight) + scale_denominator = (args.num_layers * tokens * args.moe_top_k) + scale = scale_numerator / scale_denominator + return scale * torch.dot(tokens_per_expert, expert_scores) + + +# NOTE: This class defines MoE expert computation, including expert model parallel +# communication. When using FSDP on top of MegaBlocks this is the module that should +# be wrapped s.t. the weight all-gathers can be scheduled *before* the expert model +# parallel all2all. +class ParallelMLP(torch.nn.Module): + + def __init__(self, args: Arguments): + super(ParallelMLP, self).__init__() + self.args = args + + # Calculate the number of experts in total and the number of experts + # owned by this rank. + # world_size = mpu.get_expert_parallel_world_size(args) + self.num_experts = args.moe_num_experts + self.top_k = self.args.moe_top_k + + # Calculate the number of bits needed to represent the expert indices + # so that we can pass it to radix sort. + self.sort_end_bit = max(int(np.ceil(np.log2(self.num_experts))), 1) + + # Expert MLP. + self.mlp = mlp.MLP(args) + + self.bias: Optional[torch.Tensor] + if self.args.bias: + # Note that the output bias is not parallelized with expert + # model parallelism. + self.bias = torch.nn.Parameter( + torch.empty( + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + torch.nn.init.zeros_(self.bias) + else: + self.register_parameter('bias', None) + + # Select the forward function for the operating mode. + self.forward_fn = (self.parallel_forward_once if args.moe_expert_model_parallelism else self.forward_once) + + def expert_capacity(self, tokens: int) -> int: + world_size = mpu.get_expert_parallel_world_size(self.args) + tokens_per_expert = (self.top_k * tokens * world_size / self.num_experts) + return int(self.args.moe_capacity_factor * tokens_per_expert) + + def load_balancing_loss(self, tokens_per_expert: torch.Tensor, expert_scores: torch.Tensor): + """Calculate the load balancing loss contribution.""" + assert len(expert_scores.size()) == 2 + tokens, num_experts = expert_scores.size() + assert num_experts == self.num_experts + assert len(tokens_per_expert.size()) == 1 + num_experts, = tokens_per_expert.size() + assert num_experts == self.num_experts + scale = self.num_experts / (tokens * self.top_k) + return scale * torch.dot( + tokens_per_expert.to(expert_scores.dtype), + expert_scores.mean(dim=0), + ) + + def indices_and_bins(self, + top_expert: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # Sort the expert ids to produce the scatter/gather + # indices for the permutation. + # + # TODO(tgale): Is it worth doing this conversion to 32-bit + # prior? Could we place the `torch.max` operation to return + # 32-bit expert indices? + top_expert = top_expert.int() + # output = ops.sort(top_expert, self.sort_end_bit) + output = sort(top_expert, self.sort_end_bit) + assert output is not None + bin_ids, indices = output + + # Histogram the expert ids to identify the number of + # tokens routed to each expert. + # + # TODO(tgale): Does the sorted data produce a more favorable + # data distribution for histogram? Or is the op parallelism + # worth more? + # tokens_per_expert = ops.histogram(top_expert, self.num_experts) + tokens_per_expert = histogram(top_expert, self.num_experts) + + # Calculate the bin bounds for the sorted tokens. + # bins = ops.inclusive_cumsum(tokens_per_expert, 0) + bins = inclusive_cumsum(tokens_per_expert, 0) + assert bins is not None + bins = bins.view(1) if not len(bins.size()) else bins + + assert isinstance(indices, torch.Tensor) + assert isinstance(bin_ids, torch.Tensor) + assert isinstance(bins, torch.Tensor) + assert isinstance(tokens_per_expert, torch.Tensor) + + return indices, bin_ids, bins, tokens_per_expert + + def permute_and_compute( + self, + x: torch.Tensor, + tokens_per_expert: int, # unused + indices: torch.Tensor, + bin_ids: torch.Tensor, # unused + expert_weights: torch.Tensor, + bins: torch.Tensor, + expert_capacity: int, + top_k: int, + ): + # Route the tokens for MoE computation. + x = x.view(-1, x.shape[-1]) + # output = ops.binned_gather(x, indices, bins, expert_capacity, top_k) + output = binned_gather(x, indices, bins, expert_capacity, top_k) + assert output is not None + x = output + + # Perform the expert computation. Note that we don't + # use biases for these linear operations. + x = self.mlp(x) + + # Un-route the data for the MoE output. + # return ops.binned_scatter(x, indices, expert_weights, bins, top_k) + return binned_scatter(x, indices, expert_weights, bins, top_k) + + + def forward_once(self, x: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor): + # x: [sl, bs, hs] + # expert_weights: [sl * bs, top-k] + # top_experts: [sl * bs, top-k] + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + with torch.no_grad(): + indices, bin_ids, bins, tokens_per_expert = (self.indices_and_bins(top_experts)) + + # If expert_capacity is set to zero, set the number of tokens + # per expert to the maximum we need to avoid dropping tokens. + sl, bs, _ = x.size() + expert_capacity = self.expert_capacity(sl * bs) + if expert_capacity == 0: + expert_capacity = torch.max(tokens_per_expert).item() + + x = self.permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capacity, + self.top_k, + ) + return x, tokens_per_expert + + def parallel_forward_once(self, x: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor): + # NOTE: This function implements the same computation as forward_once + # but with expert model parallelism. + # + # 1. Permute the tokens locally so that they are grouped by their + # expert assignments. This allows us to transfer all of the tokens + # for a remote device in one communication primitive. + # + # 2. Permute the tokens across the expert parallel devices. After + # this is completed each device has all of the tokens assigned to + # its set of experts in its local HBM. + # + # 3. Permute the tokens locally so that they are grouped by their + # expert assignement. After the distributed permutation the tokens + # are grouped by which device they came from. We re-order them + # locally to allow for efficient computation. + # + # After this series of permutations we compute the linear layers + # and then repeat these three steps in reverse to produce the final + # output. + # + # Compute the mapping of local tokens to experts. + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + with torch.no_grad(): + indices, bin_ids, bins, tokens_per_expert = (self.indices_and_bins(top_experts)) + + # If we're sharding the experts along the hidden dimension + # multiple devices own parts of the same sets of experts. + # Replicate the token counts so every device gets the counts. + # repeated_tokens_per_expert = ops.repeat( + repeated_tokens_per_expert = repeat( + tokens_per_expert, + (mpu.hidden_sharding_degree(self.args),), + ) + + # Pass token count information to the device on which the + # target expert resides. + parallel_tokens_per_expert = torch.empty_like(repeated_tokens_per_expert,) + tpe_handle = dist.all_to_all_single( + parallel_tokens_per_expert, + repeated_tokens_per_expert, + group=self.args.expert_parallel_group, + async_op=True, + ) + + # Permute locally and without any padding so that tokens for each + # parallel device are stored contiguously. + # + # This view updates the shape of the tensor from [sl, bs, hs] to + # [sl * bs, hs] prior to the permutation. + x = x.view(-1, x.shape[-1]) + # output = ops.gather(x, indices, bin_ids, bins, self.top_k) + output = gather(x, indices, bin_ids, bins, self.top_k) + assert output is not None + x = output + + # Compute the number of tokens that will be received from each + # device and permute the input data across the devices. + with torch.no_grad(): + tpe_handle.wait() + experts_per_rank = mpu.experts_per_rank(self.args) + + # Reshape to [world_size, num_experts_per_rank]. + world_size = mpu.get_expert_parallel_world_size(self.args) + repeated_tokens_per_expert = (repeated_tokens_per_expert.view(world_size, experts_per_rank)) + parallel_tokens_per_expert = (parallel_tokens_per_expert.view(world_size, experts_per_rank)) + + # TODO(tgale): It might be faster to do this on the GPU and + # then communicate the results back to the host. + send_counts = repeated_tokens_per_expert.cpu().sum(dim=-1) + parallel_tokens_per_expert_cpu = parallel_tokens_per_expert.cpu() + recv_counts = parallel_tokens_per_expert_cpu.sum(dim=-1) + + # Convert the send/recv counts to lists. + send_counts = send_counts.tolist() + recv_counts = recv_counts.tolist() + tokens_received = sum(recv_counts) + + # If we're sharding the experts along the hidden dimension + # multiple devices own parts of the same sets of experts. + # Replicate the token counts so devices that share experts + # get all of the tokens assigned to them. + # + # TODO(tgale): Fuse this into the prior, local permutation. + # x = ops.repeat(x, (mpu.hidden_sharding_degree(self.args), 1)) + x = repeat(x, (mpu.hidden_sharding_degree(self.args), 1)) + + # Start the cross-device permutation asynchronously so we can + # overlap communication with computation. + parallel_x, parallel_x_handle = all_to_all( + x, + recv_counts, + send_counts, + self.args.expert_parallel_group, + async_op=True, + ) + + with torch.no_grad(): + # After we do the cross-device permutation we have the tokens on the + # correct device but not yet grouped by expert because we received + # tokens from each device as contiguous chunks. To group the tokens + # for expert computation we'll do one more local permutation. The + # rest of this torch.no_grad() scope sets up the indices and bins + # for this permutation. + # replicate_bins = ops.inclusive_cumsum( + replicate_bins = inclusive_cumsum( + parallel_tokens_per_expert.flatten(), + 0, + ) + replicate_bins = (replicate_bins.view(1) if not len(replicate_bins.size()) else replicate_bins) + + # Construct the expert indices for the permuted tokens. + parallel_top_expert = torch.remainder( + torch.arange( + self.num_experts * mpu.hidden_sharding_degree(self.args), + dtype=torch.int32, + device=indices.device, + ), + mpu.experts_per_rank(self.args), + ) + # parallel_top_expert = ops.replicate( + parallel_top_expert = replicate( + parallel_top_expert.unsqueeze(dim=0), + replicate_bins, + tokens_received, + ).flatten() + + # TODO(tgale): The sort_end_bit here can be reduced. + # parallel_bin_ids, parallel_indices = ops.sort( + parallel_bin_ids, parallel_indices = sort( + parallel_top_expert, + self.sort_end_bit, + ) + + # Calculate the bins boundaries from the token counts. + parallel_tokens_per_expert = parallel_tokens_per_expert.sum( + dim=0, + dtype=torch.int, + ) + # parallel_bins = ops.inclusive_cumsum(parallel_tokens_per_expert, 0) + parallel_bins = inclusive_cumsum(parallel_tokens_per_expert, 0) + parallel_bins = (parallel_bins.view(1) if not len(parallel_bins.size()) else parallel_bins) + + # If expert_capacity is set to zero, set the number of tokens + # per expert to the maximum we need to avoid dropping tokens. + tokens, _ = x.size() + expert_capacity = self.expert_capacity(tokens) + if expert_capacity == 0: + expert_capacity = torch.max(parallel_tokens_per_expert).item() + + # Locally permute the tokens and perform the expert computation. + # Block to make sure that the cross-device permutation is complete. + if self.args.mlp_impl == 'grouped': + # GroupedMLP requires counts on CPU. We can use the tensor already + # moved to CPU for the prior all_to_all, which avoids an extra + # device synchronization. + parallel_tokens_per_expert = parallel_tokens_per_expert_cpu.sum( + dim=0, + dtype=torch.int, + ) + parallel_x_handle.wait() + parallel_x = self.permute_and_compute( + parallel_x, + parallel_tokens_per_expert, + parallel_indices, + parallel_bin_ids, + None, # expert_weights + parallel_bins, + expert_capacity, + top_k=1, + ) + + # Un-permute the tokens across the devices. + x, _ = all_to_all( + parallel_x, + send_counts, + recv_counts, + self.args.expert_parallel_group, + ) + + # Reduce along the hidden sharding to get the final outputs. + # + # TODO(tgale): Fuse this into the following local permutation. + shape = ( + mpu.hidden_sharding_degree(self.args), + -1, + self.args.hidden_size, + ) + # x = ops.sum(x.view(shape), dim=0) + x = x.view(shape).sum(dim=0) + + # Un-permute locally to setup for the next series of operations. + # x = ops.scatter(x, indices, bin_ids, expert_weights, bins, self.top_k) + x = scatter(x, indices, bin_ids, expert_weights, bins, self.top_k) + return x, tokens_per_expert.flatten() + + def forward(self, x: torch.Tensor, scores: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor): + in_shape = x.size() + + # Compute the experts. + x, tokens_per_expert = self.forward_fn(x, expert_weights, top_experts) + if self.training and self.args.moe_loss_weight > 0: + save_load_balancing_loss((tokens_per_expert, scores)) + x = x.view(in_shape) + if self.bias is not None: + if self.args.return_bias: + return x, self.bias + return x + self.bias + return x + + +class MoE(torch.nn.Module): + + def __init__(self, args: Arguments): + super(MoE, self).__init__() + + # Token router. + self.router = router.LearnedRouter(args) + + # Expert computation helper. + self.experts = self._init_experts_mlp(args) + + self.shared_expert = None + if args.shared_expert: + # SharedExpert computation helper. + self.shared_expert = sharedexpert_registry.get(args) + + def _init_experts_mlp(self, args: Arguments): + return ParallelMLP(args) + + def forward(self, x: torch.Tensor): + # NOTE: If we're going to cast the activations to lower precision + # do it before we permute the tokens to save bandwidth. + x = common.cast_if_autocast_enabled(x) + + # Compute the expert scores and assignments. + scores, expert_weights, top_experts = self.router(x) + + # Compute the experts. + out = self.experts(x, scores, expert_weights, top_experts) + if self.shared_expert is not None: + shared_expert_out = self.shared_expert(x) + out = self.shared_expert.add_experts_sharedexpert( + shared_expert_out, + out, + ) + return out diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/mpu.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/mpu.py new file mode 100644 index 0000000000000000000000000000000000000000..434e143ab42bf3f83406d69e9dd1f72777716e22 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/mpu.py @@ -0,0 +1,94 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Optional + +import torch +import torch.distributed as dist + +# from megablocks.layers.arguments import Arguments +from .arguments import Arguments + + +class MoeParam(torch.Tensor): + + def __init__(self): + super().__init__(self) + self.expert_model_parallel: bool + + +def is_moe_param(tensor: torch.Tensor) -> bool: + return hasattr(tensor, 'expert_model_parallel') + + +def get_expert_parallel_world_size(args: Arguments) -> int: + return (dist.get_world_size(args.expert_parallel_group) if args.moe_expert_model_parallelism else 1) + + +def get_expert_parallel_rank(args: Arguments) -> int: + return (dist.get_rank(args.expert_parallel_group) if args.moe_expert_model_parallelism else 0) + + +def set_expert_model_parallel_attributes( + tensor: torch.Tensor, + is_parallel: bool, +): + assert not hasattr(tensor, 'expert_model_parallel') + setattr(tensor, 'expert_model_parallel', is_parallel) + + +def param_is_expert_model_parallel(param: MoeParam) -> bool: + return (hasattr(param, 'expert_model_parallel') and param.expert_model_parallel) + + +def copy_expert_model_parallel_attributes( + destination_tensor: torch.Tensor, + source_tensor: torch.Tensor, +): + if hasattr(source_tensor, 'expert_model_parallel'): + setattr( + destination_tensor, + 'expert_model_parallel', + getattr(source_tensor, 'expert_model_parallel'), + ) + + +def synchronized_print(group: Optional[dist.ProcessGroup], *x: torch.Tensor): + world_size = dist.get_world_size(group) + rank = dist.get_rank(group) + for i in range(world_size): + dist.barrier(group) + if i == rank: + print(f'rank = {rank}', *x) + + +# Helpers for expert/tensor sharding. +def expert_sharding_degree(args: Arguments) -> int: + world_size = get_expert_parallel_world_size(args) + esd = min(world_size, args.moe_num_experts) + + if (args.moe_num_experts % esd) != 0: + raise ValueError(f'Cannot shard {args.moe_num_experts} experts {esd} ways.',) + return esd + + +def hidden_sharding_degree(args: Arguments) -> int: + world_size = get_expert_parallel_world_size(args) + esd = expert_sharding_degree(args) + hsd = world_size // esd + + if (args.ffn_hidden_size % hsd) != 0: + raise ValueError(f'Cannot shard {args.ffn_hidden_size} features {hsd} ways.',) + if (esd * hsd) != world_size: + raise ValueError( + f"Invalid sharding. 'expert_sharding_degree' ({esd}) * hidden_sharding_degree ({hsd}) != world_size ({world_size}).", + ) + return hsd + + +def experts_per_rank(args: Arguments) -> int: + return args.moe_num_experts // expert_sharding_degree(args) + + +def features_per_rank(args: Arguments) -> int: + return args.ffn_hidden_size // hidden_sharding_degree(args) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/router.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/router.py new file mode 100644 index 0000000000000000000000000000000000000000..37cb2782348d62583376f1a183c7ede83601216d --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/router.py @@ -0,0 +1,116 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch + +# from megablocks.layers import common +# from megablocks.layers.arguments import Arguments +from . import common +from .arguments import Arguments + +_ROUTER_LOGITS = [] + + +def _save_router_logits(logits: torch.Tensor, args: Arguments): + if args.moe_zloss_weight == 0: + return + global _ROUTER_LOGITS + _ROUTER_LOGITS.append(logits) + + +def clear_router_zloss(): + global _ROUTER_LOGITS + _ROUTER_LOGITS.clear() + + +def batched_router_zloss(args: Arguments): + global _ROUTER_LOGITS + + if args.moe_zloss_weight == 0: + import warnings + warnings.warn('Call to batched_router_zloss, but moe_zloss_weight=0') + return 0 + + logits_per_router = _ROUTER_LOGITS + + if args.moe_zloss_in_fp32: + logits_per_router = [logits.float() for logits in logits_per_router] + + unscaled_zloss_per_router = torch.stack([ + torch.logsumexp(logits, dim=1).square().mean() for logits in logits_per_router + ]) + + return args.moe_zloss_weight * unscaled_zloss_per_router + + +# NOTE: To enable end-to-end benchmarking without convergence we +# support a flag to force the router to assign tokens uniformly +# across the experts. We do this with a custom autograd operation +# so that PyTorch still executes the full set of router operation. +class _UniformExpertAssignment(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, num_experts: int): + out = torch.arange(x.numel(), dtype=x.dtype, device=x.device) + out = torch.remainder(out, num_experts) + return out.view(x.shape) + + +_uniform_expert_assignment = _UniformExpertAssignment.apply + + +class LearnedRouter(torch.nn.Module): + + def __init__(self, args: Arguments): + super().__init__() + self.args = args + + # Learned router parameters. + # + # NOTE: This weight matrix is not parallelized with expert model + # parallelism. Each device needs the entire router weight matrix + # so that it can route its batch of data correctly. + self.layer = torch.nn.Linear( + args.hidden_size, + args.moe_num_experts, + bias=False, + dtype=common.dtype(args), + device=args.device, + ) + args.init_method(self.layer.weight) + + def jitter(self, x: torch.Tensor): + low: float = 1.0 - self.args.moe_jitter_eps + high: float = 1.0 + self.args.moe_jitter_eps + noise = torch.rand(x.size(), dtype=x.dtype, device=x.device) + return low + noise * (high - low) + + def _top_k(self, scores: torch.Tensor): + if self.args.moe_top_k == 1: + return scores.max(dim=-1, keepdim=True) + return torch.topk(scores, self.args.moe_top_k, dim=-1) + + def forward(self, x: torch.Tensor): + if self.training and self.args.moe_jitter_eps is not None: + x = x * self.jitter(x) + + logits = self.layer(x.view(-1, x.shape[-1])) + _save_router_logits(logits, self.args) + scores = logits.softmax(dim=-1) + expert_weights, expert_indices = self._top_k(scores) + if self.args.moe_normalize_expert_weights: + expert_weights = expert_weights / torch.norm( + expert_weights, + p=self.args.moe_normalize_expert_weights, + dim=-1, + keepdim=True, + ) + + expert_indices = ( + _uniform_expert_assignment( + expert_indices, + self.args.moe_num_experts, + ) if self.args.uniform_expert_assignment else expert_indices + ) + return scores, expert_weights, expert_indices diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_layers/sharedexpert_registry.py b/build/torch214-cxx11-cu130-x86_64-linux/_layers/sharedexpert_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..5840862f88f370ace5fd49bd0612fc98d186cc49 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_layers/sharedexpert_registry.py @@ -0,0 +1,32 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Union + +# from megablocks.layers import glu, mlp +# from megablocks.layers.arguments import Arguments +from . import glu, mlp +from .arguments import Arguments + +_REGISTRY = { + 'mlp': mlp.SharedMLP, + 'glu': glu.SharedGLU, +} + + +def get(args: Arguments) -> Union[mlp.SharedMLP, glu.SharedGLU]: + """Returns an SharedMLP for use in a dMoE instance. + + Uses the provided arguments to instantiate the appropriate + SharedMLP instance. + + Args: + args: propagated Arguments dataclass. + + Returns: + An instantiated SharedMLP constructed using the input args. + """ + if args.mlp_type not in _REGISTRY: + raise ValueError(f'Unsupported mlp type: {args.mlp_type}') + + return _REGISTRY[args.mlp_type](args) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_megablocks_cuda_addf474.abi3.so b/build/torch214-cxx11-cu130-x86_64-linux/_megablocks_cuda_addf474.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..84aca1a67a2df563b6a626c60058ea6a56559594 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_megablocks_cuda_addf474.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ab71585e816db196958cf6daaafbfe55ed49c51e4ebc9b9498c00b41d583f69d +size 14587064 diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_ops.py b/build/torch214-cxx11-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..6640ac69d26eee39b52fe070e3c8571ed6a93d4e --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _megablocks_cuda_addf474 +ops = torch.ops._megablocks_cuda_addf474 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_megablocks_cuda_addf474::{op_name}" diff --git a/build/torch214-cxx11-cu130-x86_64-linux/_version.py b/build/torch214-cxx11-cu130-x86_64-linux/_version.py new file mode 100644 index 0000000000000000000000000000000000000000..c55783177af19bc03654c730c4892df8f8532279 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/_version.py @@ -0,0 +1,6 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +"""The MegaBlocks Version.""" + +__version__ = '0.11.0.dev0' diff --git a/build/torch214-cxx11-cu130-x86_64-linux/backend/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/backend/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9d4e43e9b7e6a9a1c3bde2df34914643ca5d8332 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/backend/__init__.py @@ -0,0 +1,2 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 diff --git a/build/torch214-cxx11-cu130-x86_64-linux/backend/kernels.py b/build/torch214-cxx11-cu130-x86_64-linux/backend/kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..d20139324883338ddb312e4b05a72056d74491ac --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/backend/kernels.py @@ -0,0 +1,557 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch +import triton +import triton.language as tl + +# Stub triton autotune when testing in a env that does not have CUDA +# this approach preserves the original code but enables testing without a GPU +if torch.cuda.is_available() is False: + import warnings + + warnings.warn("CUDA is not available. Triton autotuning is disabled.") + + def _no_autotune(*args, **kwargs): + def deco(fn): + return fn + return deco + + triton.autotune = _no_autotune + + +def assert_is_tensor(x, ndim): + if x.ndim != ndim: + raise ValueError(f'Expected {ndim}-tensor but got {x.ndim}-tensor') + + +def assert_is_matrix(x): + assert_is_tensor(x, 2) + + +def assert_is_vector(x): + if x.ndim != 1: + raise ValueError(f'Expected 1-tensor but got {x.ndim}-tensor') + + +def assert_equal(a, b): + if a != b: + raise ValueError(f'Expected dimensions to be equal but got {a} and {b}.',) + + +# a: (tokens, hidden_size), real. +# indices: (tokens * top_k), integer. +# bin_ids: (tokens * top_k), integer. +# weights: (tokens * top_k), real. +# bins: (num_experts), integer. +# padded_bins: (num_experts), integer. +@triton.autotune( + configs=[ + triton.Config({'BLOCK_X': 64}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=2), + triton.Config({'BLOCK_X': 256}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=4), + triton.Config({'BLOCK_X': 256}, num_warps=4), + ], + key=['NUM_COLUMNS'], +) +@triton.jit +def _padded_copy( + a, + b, + indices, + bin_ids, + weights, + bins, + padded_bins, + NUM_COLUMNS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_X: tl.constexpr, + A_TO_B: tl.constexpr, + SCALE: tl.constexpr, +): + # Our index into array 'a'. + index_a = tl.load(indices + tl.program_id(0)) + + # One threadblock per row in 'a'. Array 'b' has greater or equal + # number of rows since they could be padded. + bin_idx = tl.load(bin_ids + tl.program_id(0)) + + # Now we know what bin we're assigned to, but we need to know how + # many threadblocks were assigned to earlier bins so we can offset + # in our bin properly. + offset_in_bin = tl.program_id(0) + if bin_idx > 0: + offset_in_bin -= tl.load(bins + bin_idx - 1) + + # Load the starting index of our bin in array 'b'. + index_b = offset_in_bin + if bin_idx > 0: + index_b += tl.load(padded_bins + bin_idx - 1) + + # Offset the input and output pointers. + # + # If we're going from A to B, divide the input index to copy + # the same input repeatedly. If we're going from B to A we + # need to reduce the result. Using atomics is slow, so we + # do the reduce step in a second kernel. + offset = index_a // TOP_K if A_TO_B else index_a + a += tl.multiple_of(offset * NUM_COLUMNS, NUM_COLUMNS) + b += tl.multiple_of(index_b * NUM_COLUMNS, NUM_COLUMNS) + offsets = tl.max_contiguous(tl.arange(0, BLOCK_X), BLOCK_X) + + # Load the scale, if requested. + scale = tl.load(weights + index_a) if SCALE else 1 + + # Swap the pointers depending on the direction. + iptr = a if A_TO_B else b + optr = b if A_TO_B else a + + iterations = tl.cdiv(NUM_COLUMNS, BLOCK_X) + for _ in range(iterations): + mask = offsets < NUM_COLUMNS + x = tl.load(iptr + offsets, mask=mask) + x = x.to(tl.float32) * scale.to(tl.float32) + + tl.store(optr + offsets, x.to(optr.dtype.element_ty), mask=mask) + + offsets += BLOCK_X + + +def padded_gather(x, indices, bin_ids, weights, bins, padded_bins, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_vector(indices) + assert_is_vector(bin_ids) + assert_is_vector(bins) + assert_is_vector(padded_bins) + assert_equal(indices.shape[0], x.shape[0] * top_k) + assert_equal(bin_ids.shape[0], x.shape[0] * top_k) + assert_equal(bins.size(), padded_bins.size()) + + if weights is not None: + assert_equal(weights.shape[0], x.shape[0] * top_k) + + # NOTE: Because of the padding, the output size is dynamic. + # We load the final padded bin bound to get the output rows. + output_rows = padded_bins[-1].cpu().item() + out = torch.zeros((output_rows, x.shape[1]), dtype=x.dtype, device=x.device) + _padded_copy[(indices.shape[0],)]( + x, + out, + indices, + bin_ids, + weights, + bins, + padded_bins, + NUM_COLUMNS=x.shape[1], + A_TO_B=True, + TOP_K=top_k, + SCALE=weights is not None, + ) + return out + + +def gather(x, indices, bin_ids, weights, bins, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_vector(indices) + assert_is_vector(bin_ids) + assert_is_vector(bins) + assert_equal(indices.shape[0], x.shape[0] * top_k) + assert_equal(bin_ids.shape[0], x.shape[0] * top_k) + + if weights is not None: + assert_equal(weights.shape[0], x.shape[0] * top_k) + + # NOTE: There is no padding so the output rows equals the + # input rows multiplied by top_k. + output_rows = x.shape[0] * top_k + out = torch.empty((output_rows, x.shape[1]), dtype=x.dtype, device=x.device) + _padded_copy[(indices.shape[0],)]( + x, + out, + indices, + bin_ids, + weights, + bins, + bins, + NUM_COLUMNS=x.shape[1], + A_TO_B=True, + TOP_K=top_k, + SCALE=weights is not None, + ) + return out + + +def padded_scatter(x, indices, bin_ids, weights, bins, padded_bins, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_vector(indices) + assert_is_vector(bin_ids) + assert_is_vector(bins) + assert_is_vector(padded_bins) + assert_equal(indices.shape[0], bin_ids.shape[0]) + assert_equal(bins.size(), padded_bins.size()) + + if weights is not None: + assert_equal(indices.shape[0], weights.shape[0]) + + tokens = indices.shape[0] // top_k + out = torch.empty((tokens, top_k, x.shape[1]), dtype=x.dtype, device=x.device) + _padded_copy[(indices.shape[0],)]( + out, + x, + indices, + bin_ids, + weights, + bins, + padded_bins, + NUM_COLUMNS=x.shape[1], + A_TO_B=False, + TOP_K=top_k, + SCALE=weights is not None, + ) + + # Reduce along the top-k dimension, if needed. + return out.sum(dim=1) if top_k > 1 else out.view(tokens, x.shape[1]) + + +def scatter(x, indices, bin_ids, weights, bins, top_k): + return padded_scatter(x, indices, bin_ids, weights, bins, bins, top_k) + + +# x: (tokens, top_k, hidden_size), real +# grad: (tokens, hidden_size), real. +# wgrad: (tokens, top_k), real. +# indices: (tokens * top_k), integer. +# bin_ids: (tokens * top_k), integer. +# bins: (num_experts), integer. +# padded_bins: (num_experts), integer. +@triton.autotune( + configs=[ + triton.Config({'BLOCK_X': 64}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=2), + triton.Config({'BLOCK_X': 256}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=4), + triton.Config({'BLOCK_X': 256}, num_warps=4), + ], + key=['NUM_COLUMNS'], +) +@triton.jit +def _padded_copy_wgrad( + x, + grad, + wgrad, + indices, + bin_ids, + bins, + padded_bins, + NUM_COLUMNS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_X: tl.constexpr, +): + # Our index into 'tokens * top_k'. + index_out = tl.load(indices + tl.program_id(0)) + + # One threadblock per row in 'a'. Array 'b' has greater or equal + # number of rows since they could be padded. + bin_idx = tl.load(bin_ids + tl.program_id(0)) + + # Now we know what bin we're assigned to, but we need to know how + # many threadblocks were assigned to earlier bins so we can offset + # in our bin properly. + offset_in_bin = tl.program_id(0) + if bin_idx > 0: + offset_in_bin -= tl.load(bins + bin_idx - 1) + + # Load the starting index of our bin in array 'x'. + index_x = offset_in_bin + if bin_idx > 0: + index_x += tl.load(padded_bins + bin_idx - 1) + + # Offset the input and output pointers. + wgrad += index_out + grad += tl.multiple_of((index_out // TOP_K) * NUM_COLUMNS, NUM_COLUMNS) + x += tl.multiple_of(index_x * NUM_COLUMNS, NUM_COLUMNS) + offsets = tl.max_contiguous(tl.arange(0, BLOCK_X), BLOCK_X) + + acc = tl.zeros((BLOCK_X,), dtype=tl.float32) + iterations = tl.cdiv(NUM_COLUMNS, BLOCK_X) + for _ in range(iterations): + mask = offsets < NUM_COLUMNS + data = tl.load(x + offsets, mask=mask).to(tl.float32) + scale = tl.load(grad + offsets, mask=mask).to(tl.float32) + acc += data * scale + offsets += BLOCK_X + + # Reduce to get the final result and store. + out = tl.sum(acc).to(wgrad.dtype.element_ty) + tl.store(wgrad, out) + + +def padded_scatter_wgrad(x, grad, indices, bin_ids, bins, padded_bins, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_matrix(grad) + assert_is_vector(indices) + assert_is_vector(bin_ids) + assert_is_vector(bins) + assert_is_vector(padded_bins) + assert_equal(indices.shape[0], bin_ids.shape[0]) + assert_equal(bins.size(), padded_bins.size()) + + tokens = indices.shape[0] // top_k + out = torch.empty((tokens * top_k), dtype=x.dtype, device=x.device) + _padded_copy_wgrad[(indices.shape[0],)]( + x, + grad, + out, + indices, + bin_ids, + bins, + padded_bins, + NUM_COLUMNS=x.shape[1], + TOP_K=top_k, + ) + return out + + +def scatter_wgrad(x, grad, indices, bin_ids, bins, top_k): + return padded_scatter_wgrad(x, grad, indices, bin_ids, bins, bins, top_k) + + +# a: (tokens, hidden_size), real. +# b: (num_experts, expert_capacity, num_columns), real. +# indices: (tokens * top_k), integer. +# weights: (tokens * top_k), real. +# bins: (num_experts), integer. +@triton.autotune( + configs=[ + triton.Config({'BLOCK_X': 64}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=2), + triton.Config({'BLOCK_X': 256}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=4), + triton.Config({'BLOCK_X': 256}, num_warps=4), + ], + key=['NUM_COLUMNS'], +) +@triton.jit +def _binned_copy( + a, + b, + num_experts, + expert_capacity, + indices, + weights, + bins, + NUM_COLUMNS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_X: tl.constexpr, + A_TO_B: tl.constexpr, + SCALE: tl.constexpr, +): + # Load our indices into the output. + expert_idx = tl.program_id(0) + entry_idx = tl.program_id(1) + + # Calculate our offset into the output. + index_b = expert_idx * expert_capacity + entry_idx + + # Load the index bounds for our bin and calculate + # the number of tokens assigned to our expert. + start = 0 + if expert_idx > 0: + start = tl.load(bins + expert_idx - 1) + end = tl.load(bins + expert_idx) + num_tokens = end - start + + # Calculate our offset into the input. If we don't + # have an input exit early. + if entry_idx >= num_tokens: + return + index_a = tl.load(indices + start + entry_idx) + + # Offset the input and output pointers. + # + # If we're going from A to B, divide the input index to copy + # the same input repeatedly. If we're going from B to A we + # need to reduce the result. Using atomics is slow, so we + # do the reduce step in a second kernel. + offset = index_a // TOP_K if A_TO_B else index_a + a += tl.multiple_of(offset * NUM_COLUMNS, NUM_COLUMNS) + b += tl.multiple_of(index_b * NUM_COLUMNS, NUM_COLUMNS) + offsets = tl.max_contiguous(tl.arange(0, BLOCK_X), BLOCK_X) + + # Load the scale, if requested. + scale = tl.load(weights + index_a) if SCALE else 1 + + # Swap the pointers depending on the direction. + # + # NOTE: We need to zero the output in both directions. + iptr = a if A_TO_B else b + optr = b if A_TO_B else a + + iterations = tl.cdiv(NUM_COLUMNS, BLOCK_X) + for _ in range(iterations): + mask = offsets < NUM_COLUMNS + x = tl.load(iptr + offsets, mask=mask) + x = x.to(tl.float32) * scale.to(tl.float32) + + tl.store(optr + offsets, x.to(optr.dtype.element_ty), mask=mask) + + offsets += BLOCK_X + + +def binned_gather(x, indices, weights, bins, expert_capacity, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_vector(indices) + assert_is_vector(bins) + assert_equal(indices.shape[0], x.shape[0] * top_k) + + if weights is not None: + assert_equal(weights.shape[0], x.shape[0] * top_k) + + num_experts = bins.shape[0] + out = torch.zeros((num_experts, expert_capacity, x.shape[1]), dtype=x.dtype, device=x.device) + + _binned_copy[(num_experts, expert_capacity)]( + x, + out, + num_experts, + expert_capacity, + indices, + weights, + bins, + NUM_COLUMNS=x.shape[1], + A_TO_B=True, + TOP_K=top_k, + SCALE=weights is not None, + ) + return out + + +def binned_scatter(x, indices, weights, bins, top_k): + # Validate the input shapes. + assert_is_tensor(x, 3) + assert_is_vector(indices) + assert_is_vector(bins) + assert_equal(bins.shape[0], x.shape[0]) + + if weights is not None: + assert_equal(indices.shape[0], weights.shape[0]) + + num_experts, expert_capacity, hidden_size = x.shape + tokens = indices.shape[0] // top_k + out = torch.zeros((tokens, top_k, hidden_size), dtype=x.dtype, device=x.device) + _binned_copy[(num_experts, expert_capacity)]( + out, + x, + num_experts, + expert_capacity, + indices, + weights, + bins, + NUM_COLUMNS=hidden_size, + A_TO_B=False, + TOP_K=top_k, + SCALE=weights is not None, + ) + + # Reduce along the top-k dimension, if needed. + return out.sum(dim=1) if top_k > 1 else out.view(tokens, hidden_size) + + +# a: (tokens, hidden_size), real. +# b: (num_experts, expert_capacity, num_columns), real. +# indices: (tokens * top_k), integer. +# weights: (tokens * top_k), real. +# bins: (num_experts), integer. +@triton.autotune( + configs=[ + triton.Config({'BLOCK_X': 64}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=2), + triton.Config({'BLOCK_X': 256}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=4), + triton.Config({'BLOCK_X': 256}, num_warps=4), + ], + key=['NUM_COLUMNS'], +) +@triton.jit +def _binned_copy_wgrad( + x, + grad, + wgrad, + num_experts, + expert_capacity, + indices, + bins, + NUM_COLUMNS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_X: tl.constexpr, +): + # Load our indices into the output. + expert_idx = tl.program_id(0) + entry_idx = tl.program_id(1) + + # Calculate our offset into the output. + index_x = expert_idx * expert_capacity + entry_idx + + # Load the index bounds for our bin and calculate + # the number of tokens assigned to our expert. + start = 0 + if expert_idx > 0: + start = tl.load(bins + expert_idx - 1) + end = tl.load(bins + expert_idx) + num_tokens = end - start + + # Calculate our offset into the input. If we don't + # have an input exit early. + if entry_idx >= num_tokens: + return + index_out = tl.load(indices + start + entry_idx) + + # Offset the input and output pointers. + wgrad += index_out + grad += tl.multiple_of((index_out // TOP_K) * NUM_COLUMNS, NUM_COLUMNS) + x += tl.multiple_of(index_x * NUM_COLUMNS, NUM_COLUMNS) + offsets = tl.max_contiguous(tl.arange(0, BLOCK_X), BLOCK_X) + + acc = tl.zeros((BLOCK_X,), dtype=tl.float32) + iterations = tl.cdiv(NUM_COLUMNS, BLOCK_X) + for _ in range(iterations): + mask = offsets < NUM_COLUMNS + data = tl.load(x + offsets, mask=mask).to(tl.float32) + scale = tl.load(grad + offsets, mask=mask).to(tl.float32) + acc += data * scale + offsets += BLOCK_X + + # Reduce to get the final result and store. + out = tl.sum(acc).to(wgrad.dtype.element_ty) + tl.store(wgrad, out) + + +def binned_scatter_wgrad(x, grad, indices, bins, top_k): + # Validate the input shapes. + assert_is_tensor(x, 3) + assert_is_matrix(grad) + assert_is_vector(indices) + assert_is_vector(bins) + assert_equal(bins.shape[0], x.shape[0]) + + num_experts, expert_capacity, hidden_size = x.shape + tokens = indices.shape[0] // top_k + out = torch.zeros((tokens * top_k), dtype=x.dtype, device=x.device) + _binned_copy_wgrad[(num_experts, expert_capacity)]( + x, + grad, + out, + num_experts, + expert_capacity, + indices, + bins, + NUM_COLUMNS=hidden_size, + TOP_K=top_k, + ) + return out diff --git a/build/torch214-cxx11-cu130-x86_64-linux/benchmark_util.py b/build/torch214-cxx11-cu130-x86_64-linux/benchmark_util.py new file mode 100644 index 0000000000000000000000000000000000000000..02612d95e3ead1175a596e2878fa34b5bf85ad6f --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/benchmark_util.py @@ -0,0 +1,35 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import torch + + +def log_benchmark(name, arguments, time, std): + print('=' * 60) + print(f'{name} Benchmark') + print('Benchmark Parameters:') + for (key, value) in arguments.items(): + print(f'{key} = {value}') + print('Results:') + print('mean time = {:.3f}ms, std time = {:.3f}ms'.format(time, std)) + print('=' * 60) + + +def benchmark_function(fn, iterations=100, warmup=10): + # Warmup iterations. + for _ in range(warmup): + fn() + + times = [] + for i in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + start.record() + fn() + end.record() + + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + return np.mean(times), np.std(times) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/cpu_fused_moe.py b/build/torch214-cxx11-cu130-x86_64-linux/cpu_fused_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..7a29e81d6e9c57f64b0a78ac1c0828e45fd9d855 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/cpu_fused_moe.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# MegaBlocks CPU Fused MoE Implementation +# +# This is a pure Python/PyTorch implementation for CPU. +# For better performance, consider using the C++ kernel implementation. +# +import torch +import torch.nn.functional as F + + +def swigluoai_activation(gate: torch.Tensor, up: torch.Tensor, + alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor: + """ + SwigluOAI activation function used in GptOss models. + + Formula: + gate = clamp(gate, max=limit) + up = clamp(up, -limit, limit) + glu = gate * sigmoid(gate * alpha) + output = (up + 1) * glu + + Args: + gate: Gate tensor from gate projection + up: Up tensor from up projection + alpha: Scaling factor for sigmoid (default: 1.702) + limit: Clamp limit (default: 7.0) + + Returns: + Activated tensor + """ + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + glu = gate * torch.sigmoid(gate * alpha) + return (up + 1) * glu + + +def silu_and_mul_activation(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """ + SiLU (Swish) activation with element-wise multiplication. + + Formula: + output = silu(gate) * up + + Args: + gate: Gate tensor + up: Up tensor + + Returns: + Activated tensor + """ + return F.silu(gate) * up + + +def route_tokens_cpu( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor | None, + moe_top_k: int, + moe_num_experts: int, + moe_normalize_expert_weights: int | None = None, +) -> tuple: + """ + Route tokens to experts and compute expert weights and indices (CPU version). + + Args: + x: Input tensor [batch, seq, hidden] or [tokens, hidden] + router_weight: Router weight [num_experts, hidden] + router_bias: Router bias [num_experts] or None + moe_top_k: Number of experts per token + moe_num_experts: Total number of experts + moe_normalize_expert_weights: Normalization order or None + + Returns: + Tuple of (logits, expert_weights, expert_indices) + """ + x_flat = x.view(-1, x.shape[-1]) + logits = F.linear(x_flat, router_weight, router_bias) + + if moe_top_k == 1: + expert_weights, expert_indices = logits.max(dim=-1, keepdim=True) + else: + expert_weights, expert_indices = torch.topk(logits, moe_top_k, dim=-1) + + expert_weights = expert_weights.softmax(dim=-1) + + if moe_normalize_expert_weights is not None: + expert_weights = expert_weights / torch.norm( + expert_weights, + p=moe_normalize_expert_weights, + dim=-1, + keepdim=True, + ) + + return logits, expert_weights, expert_indices + + +def cpu_fused_moe( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + activation: str = "silu", + alpha: float = 1.702, + limit: float = 7.0, + is_interleaved: bool = True, +) -> torch.Tensor: + """ + CPU Fused MoE using PyTorch operations. + + This implementation processes all experts in parallel using batched operations + instead of sequential for loops, which is more efficient on CPU. + + Args: + hidden_states: [num_tokens, hidden_size] + w1: [num_experts, hidden_size, 2*inter_size] - gate_up_proj weights + w2: [num_experts, inter_size, hidden_size] - down_proj weights + topk_weights: [num_tokens, topk] - routing weights + topk_ids: [num_tokens, topk] - expert indices + w1_bias: [num_experts, 2*inter_size] or None + w2_bias: [num_experts, hidden_size] or None + activation: "silu" or "swigluoai" + alpha: swigluoai alpha parameter + limit: swigluoai limit parameter + is_interleaved: whether gate_up is interleaved [g0,u0,g1,u1,...] (True for GptOss) + + Returns: + output: [num_tokens, hidden_size] + """ + num_tokens, hidden_size = hidden_states.shape + num_experts = w1.shape[0] + inter_size = w2.shape[1] + topk = topk_weights.shape[1] + + # Initialize output + output = torch.zeros_like(hidden_states) + + # Build expert mask: which tokens go to which expert + # expert_mask[expert_id] contains indices of (token_idx, topk_pos) pairs + for expert_idx in range(num_experts): + # Find tokens assigned to this expert + # mask shape: [num_tokens, topk], True where topk_ids == expert_idx + mask = (topk_ids == expert_idx) + + if not mask.any(): + continue + + # Get token indices and topk positions + token_indices, topk_positions = torch.where(mask) + + if len(token_indices) == 0: + continue + + # Gather input tokens for this expert + # current_hidden: [num_selected_tokens, hidden_size] + current_hidden = hidden_states[token_indices] + + # Get weights for this expert + # w1[expert_idx]: [hidden_size, 2*inter_size] + # w2[expert_idx]: [inter_size, hidden_size] + expert_w1 = w1[expert_idx] # [hidden_size, 2*inter_size] + expert_w2 = w2[expert_idx] # [inter_size, hidden_size] + + # First projection: hidden @ w1 -> [num_selected, 2*inter_size] + gate_up = current_hidden @ expert_w1 + + # Add bias if present + if w1_bias is not None: + gate_up = gate_up + w1_bias[expert_idx] + + # Split gate and up projections + if is_interleaved: + # GptOss uses interleaved layout: [g0, u0, g1, u1, ...] + gate = gate_up[..., ::2] # [num_selected, inter_size] + up = gate_up[..., 1::2] # [num_selected, inter_size] + else: + # Standard layout: [gate_all, up_all] + gate = gate_up[..., :inter_size] + up = gate_up[..., inter_size:] + + # Apply activation + if activation == "swigluoai": + activated = swigluoai_activation(gate, up, alpha, limit) + else: # silu + activated = silu_and_mul_activation(gate, up) + + # Second projection: activated @ w2 -> [num_selected, hidden_size] + expert_out = activated @ expert_w2 + + # Add bias if present + if w2_bias is not None: + expert_out = expert_out + w2_bias[expert_idx] + + # Apply routing weights and accumulate + # weights shape: [num_selected] + weights = topk_weights[token_indices, topk_positions].unsqueeze(-1) + weighted_out = expert_out * weights + + # Accumulate to output + output.index_add_(0, token_indices, weighted_out.to(output.dtype)) + + return output + + +class MegaBlocksMoeMLP(torch.nn.Module): + """ + CPU MoE MLP module that can be used as a drop-in replacement for + the transformers GptOssMLP when using @use_kernel_forward_from_hub. + """ + can_torch_compile: bool = True + + def forward(self, x: torch.Tensor) -> tuple: + """ + Forward pass through the MoE layer. + + Args: + x: Input tensor of shape [batch_size, seq_len, hidden_size] or [tokens, hidden_size] + + Returns: + Tuple of (output, expert_weights) where: + - output: Tensor of same shape as input + - expert_weights: Expert weights for each token [tokens, top_k] + """ + # Get MoE parameters from the wrapped modules + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + moe_normalize_expert_weights = getattr( + self.experts, "normalize_expert_weights", None + ) + + # Detect activation type + if hasattr(self.experts, "alpha") and hasattr(self.experts, "limit"): + activation = "swigluoai" + alpha = self.experts.alpha + limit = self.experts.limit + else: + activation = getattr(self.experts, "activation", "silu") + alpha = 1.702 + limit = 7.0 + + # Get weight tensors + if hasattr(self.experts, "gate_up_proj"): + w1 = self.experts.gate_up_proj + is_interleaved = True # GptOss uses interleaved layout + elif hasattr(self.experts, "w1"): + w1 = self.experts.w1 + w3 = getattr(self.experts, "w3", None) + if w3 is not None: + w1 = torch.cat([w1, w3], dim=-1) + is_interleaved = False + else: + raise AttributeError("experts module must have 'gate_up_proj' or 'w1' attribute") + + if hasattr(self.experts, "down_proj"): + w2 = self.experts.down_proj + elif hasattr(self.experts, "w2"): + w2 = self.experts.w2 + else: + raise AttributeError("experts module must have 'down_proj' or 'w2' attribute") + + # Get optional bias tensors + w1_bias = getattr(self.experts, "gate_up_proj_bias", None) + w2_bias = getattr(self.experts, "down_proj_bias", None) + + # Store original shape + in_shape = x.size() + + # Route tokens to experts + logits, expert_weights, expert_indices = route_tokens_cpu( + x, + self.router.weight, + getattr(self.router, "bias", None), + moe_top_k, + moe_num_experts, + moe_normalize_expert_weights, + ) + + # Reshape input for fused MoE + x_flat = x.view(-1, x.shape[-1]) + + # Call CPU fused MoE + output = cpu_fused_moe( + hidden_states=x_flat, + w1=w1, + w2=w2, + topk_weights=expert_weights, + topk_ids=expert_indices, + w1_bias=w1_bias, + w2_bias=w2_bias, + activation=activation, + alpha=alpha, + limit=limit, + is_interleaved=is_interleaved, + ) + + # Restore original shape + output = output.view(in_shape) + + return output, expert_weights + + +# Export classes and functions +__all__ = [ + "MegaBlocksMoeMLP", + "cpu_fused_moe", + "route_tokens_cpu", + "swigluoai_activation", + "silu_and_mul_activation", +] diff --git a/build/torch214-cxx11-cu130-x86_64-linux/cpu_moe_cpp.py b/build/torch214-cxx11-cu130-x86_64-linux/cpu_moe_cpp.py new file mode 100644 index 0000000000000000000000000000000000000000..073ff66d24ce348fbb5ed19c9027fadd3f7a9c61 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/cpu_moe_cpp.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +# MegaBlocks C++ Optimized CPU MoE + +""" +C++ accelerated MoE with brgemm optimization for Intel AMX. +Direct replacement for cpu_fused_moe.MegaBlocksMoeMLP with better performance. +""" + +import torch +from typing import Optional +from .cpu_fused_moe import route_tokens_cpu +from ._ops import ops + + +def _to_local_tensor(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + """Convert DTensor to local torch.Tensor if needed for custom ops compatibility.""" + if tensor is None: + return None + # Check if it's a DTensor by looking for the to_local() method + if hasattr(tensor, "to_local"): + return tensor.to_local() + return tensor + + +def fused_moe_cpp( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + inplace: bool = False, + use_int8_w8a8: bool = False, + use_fp8_w8a16: bool = False, + use_mxfp4: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + block_size: Optional[list] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + w1_bias: Optional[torch.Tensor] = None, + w2_bias: Optional[torch.Tensor] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + is_vnni: bool = False, +) -> torch.Tensor: + """ + C++ Fused MoE with brgemm optimization (sglang compatible interface). + + Uses at::native::cpublas::brgemm for efficient batch GEMM on Intel CPUs. + Supports both silu_and_mul (standard SwiGLU) and swigluoai (GptOss) activations. + + Args: + hidden_states: Input tensor [M, K] + w1: Gate and up projections [E, 2N, K] + w2: Down projection [E, K, N] + topk_weights: Expert weights [M, topk] + topk_ids: Expert indices [M, topk] + inplace: Whether to use hidden_states as output + use_int8_w8a8: Use int8 quantization + use_fp8_w8a16: Use fp8 quantization + use_mxfp4: Use mxfp4 quantization + w1_scale, w2_scale: Quantization scales + block_size: Block size for fp8 + a1_scale, a2_scale: Activation scales + w1_bias, w2_bias: Optional biases + alpha: swigluoai alpha parameter (set to enable swiglu) + limit: swigluoai limit parameter (set to enable swiglu) + is_vnni: Whether w1/w2 are already in VNNI packed format + """ + # MXFP4/FP8 kernels only support bf16, convert if needed + orig_dtype = hidden_states.dtype + need_convert = ((use_mxfp4 or use_fp8_w8a16) and orig_dtype != torch.bfloat16) or orig_dtype == torch.float32 + if need_convert: + hidden_states = hidden_states.to(torch.bfloat16) + + # bias must match hidden_states dtype + if w1_bias is not None: + w1_bias = w1_bias.to(hidden_states.dtype) + if w2_bias is not None: + w2_bias = w2_bias.to(hidden_states.dtype) + + # Convert DTensor to local tensor for custom ops compatibility (TP mode) + hidden_states = _to_local_tensor(hidden_states) + w1 = _to_local_tensor(w1) + w2 = _to_local_tensor(w2) + topk_weights = _to_local_tensor(topk_weights) + topk_ids = _to_local_tensor(topk_ids) + w1_scale = _to_local_tensor(w1_scale) + w2_scale = _to_local_tensor(w2_scale) + a1_scale = _to_local_tensor(a1_scale) + a2_scale = _to_local_tensor(a2_scale) + w1_bias = _to_local_tensor(w1_bias) + w2_bias = _to_local_tensor(w2_bias) + + output = ops.fused_experts( + hidden_states, w1, w2, topk_weights, topk_ids, + inplace, use_int8_w8a8, use_fp8_w8a16, use_mxfp4, + w1_scale, w2_scale, block_size, a1_scale, a2_scale, + w1_bias, w2_bias, alpha, limit, is_vnni + ) + + # Convert back to original dtype if needed + if need_convert: + output = output.to(orig_dtype) + return output + + +class CPUMegaBlocksMoeMLP(torch.nn.Module): + """ + C++ optimized MoE MLP using brgemm. + Drop-in replacement for cpu_fused_moe.MegaBlocksMoeMLP with better performance. + + Usage in transformers: + # Will be used via @use_kernel_forward_from_hub decorator + """ + can_torch_compile: bool = True + + def forward(self, x: torch.Tensor) -> tuple: + """ + Forward pass through the MoE layer using C++ kernel. + + Args: + x: Input tensor [batch_size, seq_len, hidden_size] + + Returns: + Tuple of (output, expert_weights) + """ + # Optimization for GPT-OSS model + if getattr(self, "use_mxfp4", None) is None: + self.use_mxfp4 = False + + w1_scale = None + w2_scale = None + + if ( + not getattr(self, "packed_scales", False) + and hasattr(self.experts, "gate_up_proj") + and getattr(self.experts, "gate_up_proj_precision_config", None) is not None + ): + # convert scales + data_1 = ops.convert_scale_packed(self.experts.gate_up_proj_precision_config.weight_scale.data.transpose(-1, -2).contiguous()) + data_2 = ops.convert_scale_packed(self.experts.down_proj_precision_config.weight_scale.data.transpose(-1, -2).contiguous()) + self.experts.gate_up_proj_precision_config.weight_scale.storage.data = data_1 + self.experts.down_proj_precision_config.weight_scale.storage.data = data_2 + self.packed_scales = True + self.use_mxfp4 = True + + if not getattr(self, "packed_weight", False) and hasattr( + self.experts, "gate_up_proj" + ): + # convert weights + data_1 = self.experts.gate_up_proj.data.transpose(-1, -2).contiguous() + data_2 = self.experts.down_proj.data.transpose(-1, -2).contiguous() + if self.use_mxfp4: + self.experts.gate_up_proj.storage.data = ops.convert_weight_packed(data_1) + self.experts.down_proj.storage.data = ops.convert_weight_packed(data_2) + else: + # convert_weight_packed only supports bfloat16, float16, int8, fp8_e4m3 or uint8(mxfp4 or int4). + data_1 = data_1.to(torch.bfloat16) if data_1.dtype == torch.float32 else data_1 + data_2 = data_2.to(torch.bfloat16) if data_2.dtype == torch.float32 else data_2 + self.experts.gate_up_proj.data = ops.convert_weight_packed(data_1) + self.experts.down_proj.data = ops.convert_weight_packed(data_2) + + # C++ kernel does not support float32. + dtype = torch.bfloat16 if x.dtype == torch.float32 else x.dtype + if getattr(self.experts, "gate_up_proj_bias", None) is not None: + self.experts.gate_up_proj_bias.data = self.experts.gate_up_proj_bias.data.to(dtype) + if getattr(self.experts, "down_proj_bias", None) is not None: + self.experts.down_proj_bias.data = self.experts.down_proj_bias.data.to(dtype) + + self.packed_weight = True + + # Get MoE parameters + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None) + + # Detect activation type + if hasattr(self.experts, "alpha") and hasattr(self.experts, "limit"): + activation = "swigluoai" + alpha = self.experts.alpha + limit = self.experts.limit + else: + activation = getattr(self.experts, "activation", "silu") + alpha = 1.702 + limit = 7.0 + + # Get weight tensors + if hasattr(self.experts, "gate_up_proj"): + w1 = self.experts.gate_up_proj + elif hasattr(self.experts, "w1"): + w1 = self.experts.w1 + w3 = getattr(self.experts, "w3", None) + if w3 is not None: + w1 = torch.cat([w1, w3], dim=-1) + else: + raise AttributeError("experts module must have 'gate_up_proj' or 'w1' attribute") + + if hasattr(self.experts, "down_proj"): + w2 = self.experts.down_proj + elif hasattr(self.experts, "w2"): + w2 = self.experts.w2 + else: + raise AttributeError("experts module must have 'down_proj' or 'w2' attribute") + + # Get optional bias tensors + w1_bias = getattr(self.experts, "gate_up_proj_bias", None) + w2_bias = getattr(self.experts, "down_proj_bias", None) + w1_bias = w1_bias if w1_bias is None else w1_bias.data + w2_bias = w2_bias if w2_bias is None else w2_bias.data + + if self.use_mxfp4: + w1_scale = self.experts.gate_up_proj_precision_config.weight_scale.data + w2_scale = self.experts.down_proj_precision_config.weight_scale.data + + # Store original shape + in_shape = x.size() + + # Route tokens to experts (Python implementation is fast enough) + logits, expert_weights, expert_indices = route_tokens_cpu( + x, + self.router.weight, + getattr(self.router, "bias", None), + moe_top_k, + moe_num_experts, + moe_normalize_expert_weights, + ) + + # Flatten input + x_flat = x.view(-1, x.shape[-1]) + + # Determine alpha/limit for swiglu activation + use_alpha = alpha if activation == "swigluoai" else None + use_limit = limit if activation == "swigluoai" else None + + # Call C++ optimized kernel + output = fused_moe_cpp( + hidden_states=x_flat, + w1=w1.data, + w2=w2.data, + topk_weights=expert_weights, + topk_ids=expert_indices.to(torch.int32), + inplace=False, + use_int8_w8a8=False, + use_fp8_w8a16=False, + use_mxfp4=self.use_mxfp4, + w1_scale=w1_scale, + w2_scale=w2_scale, + block_size=None, + a1_scale=None, + a2_scale=None, + w1_bias=w1_bias, + w2_bias=w2_bias, + alpha=use_alpha, + limit=use_limit, + is_vnni=getattr(self, "packed_weight", False), + ) + + # Restore original shape + output = output.view(in_shape) + + return output, expert_weights + + +__all__ = ["fused_moe_cpp", "MegaBlocksMoeMLP"] diff --git a/build/torch214-cxx11-cu130-x86_64-linux/grouped_gemm/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/grouped_gemm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b91c8308f0c24f4c4171b6e4f15b6f76dabf295a --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/grouped_gemm/__init__.py @@ -0,0 +1,2 @@ +from . import ops +from . import backend diff --git a/build/torch214-cxx11-cu130-x86_64-linux/grouped_gemm/backend.py b/build/torch214-cxx11-cu130-x86_64-linux/grouped_gemm/backend.py new file mode 100644 index 0000000000000000000000000000000000000000..c7ef28ced79c830dae934177f059c1f4ddc24aad --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/grouped_gemm/backend.py @@ -0,0 +1,33 @@ +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# On ROCm there is no CUTLASS grouped GEMM; dispatch to the vendored AITER +# Triton kernels instead. On CUDA we use the compiled CUTLASS `gmm` op. +_IS_ROCM = torch.version.hip is not None + +if _IS_ROCM: + from .._grouped_gemm_triton import adapter as backend +else: + # We import the backend operations from the megablocks package as + # grouped_gemm is vendored in megablocks in this repository. + from .._ops import ops as backend # type: ignore + +def _allocate_output(a, b, batch_sizes, trans_a, trans_b): + assert not (trans_a and trans_b) + assert batch_sizes.ndim == 1, "Expected 1d tensor for batch_sizes" + assert a.ndim == 2, "Expected 2d tensor for 'a'" + assert b.ndim == (2 if trans_a else 3) + + shape = ( + (batch_sizes.shape[0], a.shape[1], b.shape[1]) + if trans_a else + (a.shape[0], (b.shape[1] if trans_b else b.shape[2])) + ) + return torch.empty(*shape, device=a.device, dtype=a.dtype) + +def gmm(a, b, batch_sizes, trans_a=False, trans_b=False, c=None): + if c is None: + c = _allocate_output(a, b, batch_sizes, trans_a, trans_b) + backend.gmm(a, b, c, batch_sizes, trans_a, trans_b) + return c diff --git a/build/torch214-cxx11-cu130-x86_64-linux/grouped_gemm/ops.py b/build/torch214-cxx11-cu130-x86_64-linux/grouped_gemm/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..4b30dd14e23837ea3b12334f4e31337ed9ad2b69 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/grouped_gemm/ops.py @@ -0,0 +1,33 @@ +from . import backend +import torch + + +class GroupedGemm(torch.autograd.Function): + + @staticmethod + def forward(ctx, a, b, batch_sizes, trans_b): + ctx.save_for_backward(a, b, batch_sizes) + ctx.trans_b = trans_b + return backend.gmm(a, b, batch_sizes, trans_a=False, trans_b=trans_b) + + @staticmethod + def backward(ctx, grad): + grad = grad.contiguous() + a, b, batch_sizes = ctx.saved_tensors + trans_b = ctx.trans_b + + agrad = None + if ctx.needs_input_grad[0]: + agrad = backend.gmm( + grad, b, batch_sizes, trans_a=False, trans_b=not trans_b) + + bgrad = None + if ctx.needs_input_grad[1]: + lhs, rhs = (grad, a) if trans_b else (a, grad) + bgrad = backend.gmm( + lhs, rhs, batch_sizes, trans_a=True, trans_b=False) + return agrad, bgrad, None, None + + +def gmm(a, b, batch_sizes, trans_b=False): + return GroupedGemm.apply(a, b, batch_sizes, trans_b) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/grouped_gemm_util.py b/build/torch214-cxx11-cu130-x86_64-linux/grouped_gemm_util.py new file mode 100644 index 0000000000000000000000000000000000000000..1d6d49fc46b0a57ad46e4179df3cc1ac2a24f7ae --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/grouped_gemm_util.py @@ -0,0 +1,31 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +import warnings + +_grouped_gemm_is_available: bool = False +try: + # import grouped_gemm + pass + _grouped_gemm_is_available = True +except ImportError as error: + warnings.warn('Grouped GEMM not available.') + + +def grouped_gemm_is_available(): + return _grouped_gemm_is_available + + +def assert_grouped_gemm_is_available(): + msg = ( + 'Grouped GEMM not available. Please run ' + '`pip install git+https://github.com/tgale96/grouped_gemm@main`.', + ) + assert _grouped_gemm_is_available, msg + + +# backend = grouped_gemm.backend if grouped_gemm_is_available() else None +# ops = grouped_gemm.ops if grouped_gemm_is_available() else None + + +from .grouped_gemm import backend as ops +from .grouped_gemm import ops as backend diff --git a/build/torch214-cxx11-cu130-x86_64-linux/layers.py b/build/torch214-cxx11-cu130-x86_64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..7c91edcd9e2d1a4ef9eac90217ff481f08ab1886 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/layers.py @@ -0,0 +1,1232 @@ +import torch +import torch.distributed as dist + +from typing import Optional, Any, TYPE_CHECKING + +from . import _layers +from . import ops + +# Conditional import for meta kernel registration +if TYPE_CHECKING: + + def register_fake(fn): + return lambda name: fn + +else: + try: + from torch.library import register_fake + except ImportError: + try: + from torch.library import impl_abstract as register_fake + except ImportError: + # Fallback for older PyTorch versions + def register_fake(op_name): + def decorator(fn): + return fn + + return decorator + + +# Meta kernel implementations for torch.compile compatibility +def _install_meta_kernels(): + """Install meta kernels for existing MegaBlocks operations""" + + # Create wrapper functions that check for compilation and return meta tensors + + # Patch ops.sort + if hasattr(ops, "sort"): + original_sort = ops.sort + + def sort_with_meta(x, end_bit=None): + if torch.compiler.is_compiling(): + # print("Using meta kernel for sort") + # Meta implementation - return tensors with correct shape/dtype/device + return torch.empty_like(x), torch.empty_like(x) + # print("Using original sort kernel") + return original_sort(x, end_bit) + + ops.sort = sort_with_meta + + # Patch ops.histogram + if hasattr(ops, "histogram"): + original_histogram = ops.histogram + + def histogram_with_meta(x, max_val): + if torch.compiler.is_compiling(): + # Meta implementation + return torch.empty((max_val,), dtype=torch.int32, device=x.device) + return original_histogram(x, max_val) + + ops.histogram = histogram_with_meta + + # Patch ops.inclusive_cumsum + if hasattr(ops, "inclusive_cumsum"): + original_inclusive_cumsum = ops.inclusive_cumsum + + def inclusive_cumsum_with_meta(x, dim): + if torch.compiler.is_compiling(): + # Meta implementation + return torch.empty_like(x) + return original_inclusive_cumsum(x, dim) + + ops.inclusive_cumsum = inclusive_cumsum_with_meta + + # Patch ops.binned_gather + if hasattr(ops, "binned_gather"): + original_binned_gather = ops.binned_gather + + def binned_gather_with_meta(x, indices, bins, bin_size, top_k): + if torch.compiler.is_compiling(): + # Meta implementation - output shape based on bin_size + if x.dim() >= 2: + hidden_size = x.size(-1) + return torch.empty( + (bin_size, x.size(1), hidden_size), + dtype=x.dtype, + device=x.device, + ) + else: + return torch.empty((bin_size,), dtype=x.dtype, device=x.device) + return original_binned_gather(x, indices, bins, bin_size, top_k) + + ops.binned_gather = binned_gather_with_meta + + # Patch ops.binned_scatter + if hasattr(ops, "binned_scatter"): + original_binned_scatter = ops.binned_scatter + + def binned_scatter_with_meta(x, indices, weights, bins, top_k): + if torch.compiler.is_compiling(): + # Meta implementation - typically reduces to 2D + if x.dim() >= 3: + return torch.empty( + (x.size(1), x.size(2)), dtype=x.dtype, device=x.device + ) + else: + return torch.empty_like(x) + return original_binned_scatter(x, indices, weights, bins, top_k) + + ops.binned_scatter = binned_scatter_with_meta + + # Patch ops.gather + if hasattr(ops, "gather"): + original_gather = ops.gather + + def gather_with_meta(x, indices, bin_ids, bins, top_k): + if torch.compiler.is_compiling(): + # Meta implementation + if x.dim() >= 2: + hidden_size = x.size(-1) + return torch.empty( + (indices.numel(), hidden_size), dtype=x.dtype, device=x.device + ) + else: + return torch.empty(indices.shape, dtype=x.dtype, device=x.device) + return original_gather(x, indices, bin_ids, bins, top_k) + + ops.gather = gather_with_meta + + # Patch ops.scatter + if hasattr(ops, "scatter"): + original_scatter = ops.scatter + + def scatter_with_meta(x, indices, bin_ids, weights, bins, top_k): + if torch.compiler.is_compiling(): + # Meta implementation - restore sequence shape + seq_len = ( + indices.size(0) // top_k + if indices.numel() > 0 and top_k > 0 + else x.size(0) + ) + if x.dim() >= 2: + return torch.empty( + (seq_len, x.size(-1)), dtype=x.dtype, device=x.device + ) + else: + return torch.empty((seq_len,), dtype=x.dtype, device=x.device) + return original_scatter(x, indices, bin_ids, weights, bins, top_k) + + ops.scatter = scatter_with_meta + + # Patch ops.replicate + if hasattr(ops, "replicate"): + original_replicate = ops.replicate + + def replicate_with_meta(x, bins, num_outputs): + if torch.compiler.is_compiling(): + # Meta implementation + return torch.empty( + (x.shape[0], num_outputs), dtype=x.dtype, device=x.device + ) + return original_replicate(x, bins, num_outputs) + + ops.replicate = replicate_with_meta + + # Patch ops.repeat (if it's a regular function) + if hasattr(ops, "repeat"): + original_repeat = ops.repeat + + def repeat_with_meta(x, repeats): + if torch.compiler.is_compiling(): + # Meta implementation + if isinstance(repeats, (tuple, list)): + new_shape = list(x.shape) + for i, rep in enumerate(repeats): + if i < len(new_shape): + new_shape[i] *= rep + return torch.empty(new_shape, dtype=x.dtype, device=x.device) + else: + new_shape = [x.size(0) * repeats] + list(x.shape[1:]) + return torch.empty(new_shape, dtype=x.dtype, device=x.device) + return original_repeat(x, repeats) + + ops.repeat = repeat_with_meta + + +# Install meta kernels on import +try: + _install_meta_kernels() +except Exception as e: + # If meta kernel installation fails, continue without them + # torch.compile may not work but the library will still function + import warnings + + warnings.warn( + f"Failed to install meta kernels for torch.compile support: {e}", UserWarning + ) + + +# Set the expert model parallel attributes on a tensor +def set_expert_model_parallel_attributes( + tensor: torch.Tensor, + is_parallel: bool, +): + assert not hasattr(tensor, "expert_model_parallel") + setattr(tensor, "expert_model_parallel", is_parallel) + + +# Get the expert model parallel attributes from a tensor +def expert_sharding_degree( + world_size: int, + moe_num_experts: int, +) -> int: + esd = min(world_size, moe_num_experts) + if (moe_num_experts % esd) != 0: + raise ValueError(f"Cannot shard {moe_num_experts} experts {esd} ways.") + return esd + + +# Calculate the hidden sharding degree based on world size and expert sharding degree +def hidden_sharding_degree( + world_size: int, + moe_num_experts: int, + ffn_hidden_size: int, +) -> int: + esd = expert_sharding_degree(world_size, moe_num_experts) + hsd = world_size // esd + if (ffn_hidden_size % hsd) != 0: + raise ValueError(f"Cannot shard {ffn_hidden_size} features {hsd} ways.") + if (esd * hsd) != world_size: + raise ValueError( + f"Invalid sharding. expert_sharding_degree ({esd}) * hidden_sharding_degree ({hsd}) != world_size ({world_size})." + ) + return hsd + + +# Calculate the number of experts per rank based on world size and expert sharding degree +def experts_per_rank( + moe_num_experts: int, + world_size: int, +) -> int: + return moe_num_experts // expert_sharding_degree(world_size, moe_num_experts) + + +# Calculate the number of features per rank based on ffn hidden size and hidden sharding degree +def features_per_rank( + ffn_hidden_size: int, world_size: int, moe_num_experts: int +) -> int: + return ffn_hidden_size // hidden_sharding_degree( + world_size, moe_num_experts, ffn_hidden_size + ) + + +# Apply jitter to the input tensor +def apply_jitter(x: torch.Tensor, moe_jitter_eps: float) -> torch.Tensor: + low = 1.0 - moe_jitter_eps + high = 1.0 + moe_jitter_eps + noise = torch.rand(x.size(), dtype=x.dtype, device=x.device) + return x * (low + noise * (high - low)) + + +# Compute the top-k scores from the logits +def compute_top_k(scores: torch.Tensor, moe_top_k: int): + if moe_top_k == 1: + return scores.max(dim=-1, keepdim=True) + return torch.topk(scores, moe_top_k, dim=-1) + + +# Route tokens to experts and compute expert weights and indices +def route_tokens( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + moe_top_k: int, + moe_num_experts: int, + moe_jitter_eps: float = None, + moe_normalize_expert_weights: int = None, + uniform_expert_assignment: bool = False, + training: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if training and moe_jitter_eps is not None: + x = apply_jitter(x, moe_jitter_eps) + + x_flat = x.view(-1, x.shape[-1]) + logits = torch.nn.functional.linear(x_flat, router_weight, router_bias) + expert_weights, expert_indices = compute_top_k(logits, moe_top_k) + expert_weights = expert_weights.softmax(dim=-1) + if moe_normalize_expert_weights is not None: + expert_weights = expert_weights / torch.norm( + expert_weights, + p=moe_normalize_expert_weights, + dim=-1, + keepdim=True, + ) + if uniform_expert_assignment: + expert_indices = _layers.router._uniform_expert_assignment( + expert_indices, + moe_num_experts, + ) + + return logits, expert_weights, expert_indices + + +# Scale the gradient of the weights +def scale_grad( + w: torch.Tensor, + gradient_scale: Optional[float] = None, +) -> torch.Tensor: + if gradient_scale is None: + return w + return _layers.mlp.scale_gradient(w, gradient_scale) + + +# Forward pass for the MLP layer +def mlp_forward( + x: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w1_bias: torch.Tensor, + w2_bias: torch.Tensor, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + limit: float = 7.0, +): + # Scale weights + w1 = scale_grad(w1, gradient_scale) + w2 = scale_grad(w2, gradient_scale) + w1_bias = scale_grad(w1_bias, gradient_scale) + w2_bias = scale_grad(w2_bias, gradient_scale) + + # Resolve dtensors + w1 = _layers.mlp.resolve_dtensor(w1) + w2 = _layers.mlp.resolve_dtensor(w2) + w1_bias = _layers.mlp.resolve_dtensor(w1_bias) + w2_bias = _layers.mlp.resolve_dtensor(w2_bias) + + # Forward pass + gate_up = torch.bmm(x, w1) + w1_bias[..., None, :] + gate, up = gate_up[..., ::2], gate_up[..., 1::2] + gate = gate.clamp(min=None, max=limit) + up = up.clamp(min=-limit, max=limit) + glu = gate * torch.sigmoid(gate * alpha) + next_states = torch.bmm(((up + 1) * glu), w2) + next_states += w2_bias[..., None, :] + return next_states + +# Shared expert MLP forward pass +def shared_mlp_forward( + x: torch.Tensor, + up_proj_weight: torch.Tensor, + down_proj_weight: torch.Tensor, + up_proj_bias: Optional[torch.Tensor] = None, + down_proj_bias: Optional[torch.Tensor] = None, + activation_fn: Optional[Any] = None, + gradient_scale: Optional[float] = None, +) -> torch.Tensor: + # Default activation function + if activation_fn is None: + activation_fn = torch.nn.functional.gelu + + # Scale weights + up_proj_weight = scale_grad(up_proj_weight, gradient_scale) + down_proj_weight = scale_grad(down_proj_weight, gradient_scale) + if up_proj_bias is not None: + up_proj_bias = scale_grad(up_proj_bias, gradient_scale) + if down_proj_bias is not None: + down_proj_bias = scale_grad(down_proj_bias, gradient_scale) + + # Resolve dtensors + up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight) + down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight) + if up_proj_bias is not None: + up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias) + if down_proj_bias is not None: + down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias) + + # Up projection + x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias) + + # Activation + x = activation_fn(x) + + # Down projection + x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias) + + return x + + +# Combine outputs from shared expert and regular experts +def combine_expert_shared_outputs( + shared_expert_out: torch.Tensor, + expert_out: torch.Tensor, + shared_expert_weighted_sum: bool = False, + moe_top_k: int = 1, +) -> torch.Tensor: + if shared_expert_weighted_sum: + # Weighted sum based on number of experts used + total_experts = moe_top_k + 1 + shared_weight = 1.0 / total_experts + expert_weight = moe_top_k / total_experts + return shared_expert_out * shared_weight + expert_out * expert_weight + else: + # Simple addition + return shared_expert_out + expert_out + + +# Global variable to store load balancing loss +_LOAD_BALANCING_LOSS = [] + + +def save_load_balancing_loss(loss): + global _LOAD_BALANCING_LOSS + _LOAD_BALANCING_LOSS.append(loss) + + +def get_load_balancing_loss(): + global _LOAD_BALANCING_LOSS + return _LOAD_BALANCING_LOSS + + +def clear_load_balancing_loss(): + global _LOAD_BALANCING_LOSS + _LOAD_BALANCING_LOSS.clear() + + +def batched_load_balancing_loss(args): + if args.moe_loss_weight == 0: + return 0.0 + + tokens_per_expert, expert_scores = zip(*get_load_balancing_loss()) + num_layers_per_pipeline_stage = args.num_layers // args.pipeline_model_parallel_size + if args.num_layers_per_virtual_pipeline_stage is not None: + num_layers_per_pipeline_stage = args.num_layers_per_virtual_pipeline_stage + + if len(tokens_per_expert) != num_layers_per_pipeline_stage: + raise ValueError( + f"Expected {num_layers_per_pipeline_stage} token_per_experts " + f"but found {len(tokens_per_expert)}.\nnum_layers = " + f"{args.num_layers}\npipeline_model_parallel_size = " + f"{args.pipeline_model_parallel_size}\n" + "num_layers_per_virtual_pipeline_stage" + f" = {args.num_layers_per_virtual_pipeline_stage}", + ) + if len(expert_scores) != num_layers_per_pipeline_stage: + raise ValueError( + f"Expected {num_layers_per_pipeline_stage} expert_scores " + f"but found {len(tokens_per_expert)}.\nnum_layers = " + f"{args.num_layers}\npipeline_model_parallel_size = " + f"{args.pipeline_model_parallel_size}\n" + "num_layers_per_virtual_pipeline_stage" + f" = {args.num_layers_per_virtual_pipeline_stage}", + ) + + # Verify the shape of the tokens_per_expert and expert_scores tensors. + assert all( + (x.ndim == 1 and x.numel() == args.moe_num_experts for x in tokens_per_expert) + ) + + tokens = expert_scores[0].shape[0] + assert all( + ( + ( + x.ndim == 2 + and x.shape[1] == args.moe_num_experts + and x.shape[0] == tokens + ) + for x in expert_scores + ) + ) + + # Concatenate the contributions of each layer and convert to + # the correct types and formats for the dot product. + expert_scores = torch.cat(expert_scores, dim=1) + if args.moe_lbl_in_fp32: + expert_scores = expert_scores.float() + if tokens != 0: + expert_scores = expert_scores.mean(dim=0) + else: + expert_scores = expert_scores.sum(dim=0) + tokens_per_expert = torch.cat(tokens_per_expert).to(expert_scores.dtype) + + expected_values = num_layers_per_pipeline_stage * args.moe_num_experts + assert tokens_per_expert.numel() == expected_values + assert expert_scores.numel() == expected_values + + # Calculate the total scale across all factors. + # + # loss_weight * num_experts / (num_layers * tokens * top_k) + scale_numerator = args.moe_num_experts * args.moe_loss_weight + scale_denominator = args.num_layers * tokens * args.moe_top_k + scale = scale_numerator / scale_denominator + return scale * torch.dot(tokens_per_expert, expert_scores) + + +# Calculate the expert capacity based on tokens, top_k, number of experts, +# expert parallel group, capacity factor, and whether expert model parallelism is used. +def expert_capacity( + tokens: int, + top_k: int, + num_experts: int, + expert_parallel_group: int, + moe_capacity_factor: float, + moe_expert_model_parallelism: bool, +) -> int: + world_size = ( + dist.get_world_size(expert_parallel_group) + if moe_expert_model_parallelism + else 1 + ) + + tokens_per_expert = top_k * tokens * world_size / num_experts + return int(moe_capacity_factor * tokens_per_expert) + + +def load_balancing_loss( + tokens_per_expert: torch.Tensor, + expert_scores: torch.Tensor, + top_k: int, + num_experts: int, +): + assert len(expert_scores.size()) == 2 + tokens, num_experts = expert_scores.size() + assert num_experts == num_experts + assert len(tokens_per_expert.size()) == 1 + (num_experts,) = tokens_per_expert.size() + assert num_experts == num_experts + scale = num_experts / (tokens * top_k) + return scale * torch.dot( + tokens_per_expert.to(expert_scores.dtype), + expert_scores.mean(dim=0), + ) + + +def indices_and_bins( + top_expert: torch.Tensor, + sort_end_bit: int, + num_experts: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + top_expert = top_expert.int() + + # Ensure contiguous memory layout + top_expert = top_expert.contiguous() + + # Ensure CUB knows which device to use + with torch.cuda.device(top_expert.device): + output = ops.sort(top_expert, sort_end_bit) + bin_ids, indices = output + tokens_per_expert = ops.histogram(top_expert, num_experts) + bins = ops.inclusive_cumsum(tokens_per_expert, 0) + + bins = bins.view(1) if not len(bins.size()) else bins + return indices, bin_ids, bins, tokens_per_expert + + +def expert_capacity_fn( + tokens: int, + top_k: int, + num_experts: int, + expert_parallel_group: torch.distributed.ProcessGroup, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = False, +) -> int: + world_size = ( + dist.get_world_size(expert_parallel_group) + if moe_expert_model_parallelism + else 1 + ) + tokens_per_expert = top_k * tokens * world_size / num_experts + return int(moe_capacity_factor * tokens_per_expert) + + +def permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capacity, + top_k, + w1, + w2, + w1_bias, + w2_bias, + gradient_scale, + alpha, +): + # Route tokens to experts + x = x.view(-1, x.shape[-1]) + + # Ensure CUB knows which device to use + with torch.cuda.device(x.device): + x = ops.binned_gather(x, indices, bins, expert_capacity, top_k) + + # Expert computation + x = mlp_forward(x, w1, w2, w1_bias, w2_bias, gradient_scale, alpha) + + # Ensure CUB knows which device to use + with torch.cuda.device(x.device): + # Route tokens back + out = ops.binned_scatter(x, indices, expert_weights, bins, top_k) + return out + + +def forward_once( + x: torch.Tensor, + expert_weights: torch.Tensor, + top_experts: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w1_bias: torch.Tensor, + w2_bias: torch.Tensor, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + sort_end_bit: int = 0, + top_k: int = 4, + num_experts: int = 128, + expert_parallel_group: int = None, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = False, + mlp_impl: Optional[str] = None, +): + # x: [sl, bs, hs] + # expert_weights: [sl * bs, top-k] + # top_experts: [sl * bs, top-k] + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + + with torch.no_grad(): + indices, bin_ids, bins, tokens_per_expert = indices_and_bins( + top_experts, sort_end_bit, num_experts + ) + + # Calculate expert capacity + sl, bs, _ = x.size() + + expert_capacity = expert_capacity_fn( + sl * bs, + top_k, + num_experts, + expert_parallel_group, + moe_capacity_factor, + moe_expert_model_parallelism, + ) + + if expert_capacity == 0: + expert_capacity = torch.max(tokens_per_expert).item() + + x = permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capacity, + top_k, + w1, + w2, + w1_bias, + w2_bias, + gradient_scale, + alpha, + ) + return x, tokens_per_expert + + +def parallel_forward_once( + x: torch.Tensor, + expert_weights: torch.Tensor, + top_experts: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w1_bias: torch.Tensor, + w2_bias: torch.Tensor, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + sort_end_bit: int = 0, + top_k: int = 4, + num_experts: int = 128, + expert_parallel_group: torch.distributed.ProcessGroup = None, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = True, + hidden_size: int = 1152, + mlp_impl: Optional[str] = "grouped", +): + # Flatten inputs + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + + # TODO: remove debugging var + # my_rank = dist.get_rank(expert_parallel_group) if expert_parallel_group else 0 + + with torch.no_grad(): + # Step 1: Local permutation setup + indices, bin_ids, bins, tokens_per_expert = indices_and_bins( + top_experts, sort_end_bit, num_experts + ) + + # Calculate sharding parameters + world_size = dist.get_world_size(expert_parallel_group) + hidden_sharding_deg = hidden_sharding_degree( + world_size, num_experts, hidden_size + ) + experts_per_rank_val = experts_per_rank(num_experts, world_size) + + # Replicate token counts for hidden sharding + repeated_tokens_per_expert = ops.repeat( + tokens_per_expert, (hidden_sharding_deg,) + ) + + # Exchange token counts across devices + parallel_tokens_per_expert = torch.empty_like(repeated_tokens_per_expert) + + # Ensure CUB knows which device to use + tpe_handle = dist.all_to_all_single( + parallel_tokens_per_expert, + repeated_tokens_per_expert, + group=expert_parallel_group, + async_op=True, + ) + + # Step 2: Local permutation - group tokens by target device + x = x.view(-1, x.shape[-1]) # [sl * bs, hs] + x = ops.gather(x, indices, bin_ids, bins, top_k) + + # Step 3: Compute communication counts and exchange tokens + with torch.no_grad(): + tpe_handle.wait() + + # Reshape for per-device calculations + repeated_tokens_per_expert = repeated_tokens_per_expert.view( + world_size, experts_per_rank_val + ) + parallel_tokens_per_expert = parallel_tokens_per_expert.view( + world_size, experts_per_rank_val + ) + + # Calculate send/recv counts + send_counts = repeated_tokens_per_expert.cpu().sum(dim=-1).tolist() + # recv_counts = parallel_tokens_per_expert.cpu().sum(dim=-1).tolist() + parallel_tokens_per_expert_cpu = parallel_tokens_per_expert.cpu() + recv_counts = parallel_tokens_per_expert_cpu.sum(dim=-1).tolist() + tokens_received = sum(recv_counts) + + # Replicate for hidden sharding + x = ops.repeat(x, (hidden_sharding_deg, 1)) + + # Cross-device token exchange + parallel_x, parallel_x_handle = _layers.all_to_all.all_to_all( + x, recv_counts, send_counts, expert_parallel_group, async_op=True + ) + + with torch.no_grad(): + # Step 4: Setup for local expert computation + replicate_bins = ops.inclusive_cumsum(parallel_tokens_per_expert.flatten(), 0) + replicate_bins = ( + replicate_bins.view(1) if not len(replicate_bins.size()) else replicate_bins + ) + + # Create expert indices for received tokens + parallel_top_expert = torch.remainder( + torch.arange( + num_experts * hidden_sharding_deg, + dtype=torch.int32, + device=indices.device, + ), + experts_per_rank_val, + ) + parallel_top_expert = ops.replicate( + parallel_top_expert.unsqueeze(dim=0), + replicate_bins, + tokens_received, + ).flatten() + + # Sort tokens by expert assignment + parallel_bin_ids, parallel_indices = ops.sort( + parallel_top_expert, + sort_end_bit, + ) + + # Calculate bins for local experts + parallel_tokens_per_expert = parallel_tokens_per_expert.sum( + dim=0, dtype=torch.int + ) + parallel_bins = ops.inclusive_cumsum(parallel_tokens_per_expert, 0) + parallel_bins = ( + parallel_bins.view(1) if not len(parallel_bins.size()) else parallel_bins + ) + + # Calculate expert capacity + expert_capacity = expert_capacity_fn( + tokens_received, + top_k, + experts_per_rank_val, + expert_parallel_group, + moe_capacity_factor, + moe_expert_model_parallelism, + ) + if expert_capacity == 0: + expert_capacity = torch.max(parallel_tokens_per_expert).item() + + # Locally permute the tokens and perform the expert computation. + # Block to make sure that the cross-device permutation is complete. + if mlp_impl == "grouped": + # GroupedMLP requires counts on CPU. We can use the tensor already + # moved to CPU for the prior all_to_all, which avoids an extra + # device synchronization. + parallel_tokens_per_expert = parallel_tokens_per_expert_cpu.sum( + dim=0, + dtype=torch.int, + ) + + # Step 5: Expert computation + parallel_x_handle.wait() + + parallel_x = permute_and_compute( + parallel_x, + parallel_tokens_per_expert, + parallel_indices, + parallel_bin_ids, + None, # expert_weights + parallel_bins, + expert_capacity, + top_k=1, + w1=w1, + w2=w2, + w1_bias=w1_bias, + w2_bias=w2_bias, + gradient_scale=gradient_scale, + alpha=alpha, + ) + + # Step 6: Reverse communication - send results back + x, _ = _layers.all_to_all.all_to_all( + parallel_x, send_counts, recv_counts, expert_parallel_group + ) + + # Step 7: Reduce across hidden sharding dimension + shape = (hidden_sharding_deg, -1, hidden_size) + x = x.view(shape).sum(dim=0) + + # Step 8: Final local unpermutation + x = ops.scatter(x, indices, bin_ids, expert_weights, bins, top_k) + + return x, tokens_per_expert.flatten() + + +def moe_forward( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: Optional[torch.Tensor], + moe_top_k: int, + moe_num_experts: int, + moe_jitter_eps: float = None, + moe_normalize_expert_weights: int = None, + uniform_expert_assignment: bool = False, + training: bool = False, + w1: torch.Tensor = None, + w2: torch.Tensor = None, + w1_bias: torch.Tensor = None, + w2_bias: torch.Tensor = None, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + sort_end_bit: int = 0, + expert_parallel_group: torch.distributed.ProcessGroup = None, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = False, + forward_fn: Any = None, + hidden_size: int = None, + mlp_impl: str = "grouped", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + # Route tokens to experts + logits, expert_weights, expert_indices = route_tokens( + x, + router_weight, + router_bias, + moe_top_k, + moe_num_experts, + moe_jitter_eps, + moe_normalize_expert_weights, + uniform_expert_assignment, + training, + ) + + # Create router scores for output + router_scores = ( + torch.zeros_like(logits) + .scatter_(1, expert_indices, expert_weights) + .transpose(0, 1) + ) + + in_shape = x.size() + + # Prepare forward function arguments + forward_args = { + "x": x, + "expert_weights": expert_weights, + "top_experts": expert_indices, + "w1": w1, + "w2": w2, + "w1_bias": w1_bias, + "w2_bias": w2_bias, + "gradient_scale": gradient_scale, + "alpha": alpha, + "sort_end_bit": sort_end_bit, + "top_k": moe_top_k, + "num_experts": moe_num_experts, + "expert_parallel_group": expert_parallel_group, + "moe_capacity_factor": moe_capacity_factor, + "moe_expert_model_parallelism": moe_expert_model_parallelism, + "mlp_impl": mlp_impl, + } + + # Add hidden_size for parallel forward + if moe_expert_model_parallelism and hidden_size is not None: + forward_args["hidden_size"] = hidden_size + elif moe_expert_model_parallelism and hidden_size is None: + # Infer hidden_size from input shape + forward_args["hidden_size"] = x.shape[-1] + + # Compute expert outputs + x, tokens_per_expert = forward_fn(**forward_args) + + # Save load balancing loss if needed + moe_loss_weight = 0.0 # Can be made configurable + if training and moe_loss_weight > 0: + save_load_balancing_loss((tokens_per_expert, logits)) + + # Restore original shape + x = x.view(in_shape) + + return x, expert_weights, router_scores + + +def moe_forward_with_shared_expert( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: Optional[torch.Tensor], + moe_top_k: int, + moe_num_experts: int, + moe_jitter_eps: float = None, + moe_normalize_expert_weights: int = None, + uniform_expert_assignment: bool = False, + training: bool = False, + w1: torch.Tensor = None, + w2: torch.Tensor = None, + w1_bias: torch.Tensor = None, + w2_bias: torch.Tensor = None, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + sort_end_bit: int = 0, + expert_parallel_group: torch.distributed.ProcessGroup = None, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = False, + forward_fn: Any = None, + hidden_size: int = None, + mlp_impl: str = "grouped", + # Shared expert parameters + shared_up_proj_weight: Optional[torch.Tensor] = None, + shared_down_proj_weight: Optional[torch.Tensor] = None, + shared_up_proj_bias: Optional[torch.Tensor] = None, + shared_down_proj_bias: Optional[torch.Tensor] = None, + shared_expert_weighted_sum: bool = False, + shared_activation_fn: Optional[Any] = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + # First, compute regular MoE forward pass + expert_out, expert_weights, router_scores = moe_forward( + x=x, + router_weight=router_weight, + router_bias=router_bias, + moe_top_k=moe_top_k, + moe_num_experts=moe_num_experts, + moe_jitter_eps=moe_jitter_eps, + moe_normalize_expert_weights=moe_normalize_expert_weights, + uniform_expert_assignment=uniform_expert_assignment, + training=training, + w1=w1, + w2=w2, + w1_bias=w1_bias, + w2_bias=w2_bias, + gradient_scale=gradient_scale, + alpha=alpha, + sort_end_bit=sort_end_bit, + expert_parallel_group=expert_parallel_group, + moe_capacity_factor=moe_capacity_factor, + moe_expert_model_parallelism=moe_expert_model_parallelism, + forward_fn=forward_fn, + hidden_size=hidden_size, + mlp_impl=mlp_impl, + ) + + # If shared expert weights provided, compute shared expert output + if shared_up_proj_weight is not None and shared_down_proj_weight is not None: + shared_expert_out = shared_mlp_forward( + x=x, + up_proj_weight=shared_up_proj_weight, + down_proj_weight=shared_down_proj_weight, + up_proj_bias=shared_up_proj_bias, + down_proj_bias=shared_down_proj_bias, + activation_fn=shared_activation_fn, + gradient_scale=gradient_scale, + ) + + # Combine expert outputs + combined_out = combine_expert_shared_outputs( + shared_expert_out=shared_expert_out, + expert_out=expert_out, + shared_expert_weighted_sum=shared_expert_weighted_sum, + moe_top_k=moe_top_k, + ) + + return combined_out, expert_weights, router_scores + + # Return regular MoE output if no shared expert + return expert_out, expert_weights, router_scores + + +def create_shared_expert_weights( + hidden_size: int, + shared_expert_hidden_size: int, + device: torch.device, + dtype: torch.dtype, + init_method: Any, + output_layer_init_method: Any = None, +) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + + if output_layer_init_method is None: + output_layer_init_method = init_method + + # Create weight tensors + up_proj_weight = torch.empty( + shared_expert_hidden_size, + hidden_size, + device=device, + dtype=dtype, + ) + down_proj_weight = torch.empty( + hidden_size, + shared_expert_hidden_size, + device=device, + dtype=dtype, + ) + + # Initialize weights + init_method(up_proj_weight) + output_layer_init_method(down_proj_weight) + + # No bias by default + return up_proj_weight, down_proj_weight, None, None + + +# HACK: Extract device_mesh from pre-hook closure - required for transformers integration +# This exists because device_mesh is trapped in hook closures with no model attribute +# Fragile - breaks if hook structure changes or Python internals change +# TODO: Replace with a more robust solution when available +def get_device_mesh(model): + # Extract device_mesh from child's unused pre_hook closure + try: + # Find the pre-hook that contains 'device_mesh' in its closure + hook = next( + h + for h in model.experts._forward_pre_hooks.values() + if "device_mesh" in h.__code__.co_freevars + ) + # Extract the device_mesh from the closure + return hook.__closure__[ + hook.__code__.co_freevars.index("device_mesh") + ].cell_contents + except Exception: + return None + + +class MegaBlocksMoeMLP(torch.nn.Module): + can_torch_compile: bool = True + + def forward(self, x: torch.Tensor) -> torch.Tensor: + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + gradient_scale = getattr(self.experts, "gradient_scale", None) + alpha = getattr(self.experts, "alpha", 1.0) + moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0) + moe_jitter_eps = getattr(self.experts, "jitter_eps", None) + moe_normalize_expert_weights = getattr( + self.experts, "normalize_expert_weights", None + ) + uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False) + + expert_parallel_group = getattr(self, "expert_parallel_group", None) + if expert_parallel_group is None: + device_mesh = get_device_mesh(self) + expert_parallel_group = device_mesh.get_group() if device_mesh else None + + has_parallel = ( + expert_parallel_group is not None + and dist.is_initialized() + and dist.get_world_size(expert_parallel_group) > 1 + ) + forward_fn = parallel_forward_once if has_parallel else forward_once + + sort_end_bit = max( + int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1 + ) + mlp_impl = getattr(self, "mlp_impl", "grouped") + output, expert_weights_out, *_ = moe_forward( + x=x, + router_weight=self.router.weight, + router_bias=self.router.bias, + moe_top_k=moe_top_k, + moe_num_experts=moe_num_experts, + moe_jitter_eps=moe_jitter_eps, + moe_normalize_expert_weights=moe_normalize_expert_weights, + uniform_expert_assignment=uniform_expert_assignment, + training=self.training, + w1=self.experts.gate_up_proj, + w2=self.experts.down_proj, + w1_bias=self.experts.gate_up_proj_bias, + w2_bias=self.experts.down_proj_bias, + gradient_scale=gradient_scale, + alpha=alpha, + sort_end_bit=sort_end_bit, + expert_parallel_group=expert_parallel_group, + moe_capacity_factor=moe_capacity_factor, + moe_expert_model_parallelism=has_parallel, + forward_fn=forward_fn, + hidden_size=self.experts.hidden_size, + mlp_impl=mlp_impl, + ) + return output, expert_weights_out + + +# Export main classes +__all__ = ["MegaBlocksMoeMLP", "MegaBlocksMoeMLPWithSharedExpert"] + + +class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP): + + def __init__(self): + super().__init__() + # Shared expert weights will be set by the user + self.shared_up_proj_weight = None + self.shared_down_proj_weight = None + self.shared_up_proj_bias = None + self.shared_down_proj_bias = None + self.shared_expert_weighted_sum = False + self.shared_activation_fn = None + + def set_shared_expert_weights( + self, + up_proj_weight: torch.Tensor, + down_proj_weight: torch.Tensor, + up_proj_bias: Optional[torch.Tensor] = None, + down_proj_bias: Optional[torch.Tensor] = None, + weighted_sum: bool = False, + activation_fn: Optional[Any] = None, + ): + self.shared_up_proj_weight = up_proj_weight + self.shared_down_proj_weight = down_proj_weight + self.shared_up_proj_bias = up_proj_bias + self.shared_down_proj_bias = down_proj_bias + self.shared_expert_weighted_sum = weighted_sum + self.shared_activation_fn = activation_fn + + def forward(self, x: torch.Tensor) -> torch.Tensor: + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + gradient_scale = getattr(self.experts, "gradient_scale", None) + alpha = getattr(self.experts, "alpha", 1.0) + moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0) + moe_jitter_eps = getattr(self.experts, "jitter_eps", None) + moe_normalize_expert_weights = getattr( + self.experts, "normalize_expert_weights", None + ) + uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False) + + expert_parallel_group = getattr(self, "expert_parallel_group", None) + if expert_parallel_group is None: + device_mesh = get_device_mesh(self) + expert_parallel_group = device_mesh.get_group() if device_mesh else None + + has_parallel = ( + expert_parallel_group is not None + and dist.is_initialized() + and dist.get_world_size(expert_parallel_group) > 1 + ) + forward_fn = parallel_forward_once if has_parallel else forward_once + + sort_end_bit = max( + int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1 + ) + mlp_impl = getattr(self, "mlp_impl", "grouped") + + output, expert_weights_out, *_ = moe_forward_with_shared_expert( + x=x, + router_weight=self.router.weight, + router_bias=self.router.bias, + moe_top_k=moe_top_k, + moe_num_experts=moe_num_experts, + moe_jitter_eps=moe_jitter_eps, + moe_normalize_expert_weights=moe_normalize_expert_weights, + uniform_expert_assignment=uniform_expert_assignment, + training=self.training, + w1=self.experts.gate_up_proj, + w2=self.experts.down_proj, + w1_bias=self.experts.gate_up_proj_bias, + w2_bias=self.experts.down_proj_bias, + gradient_scale=gradient_scale, + alpha=alpha, + sort_end_bit=sort_end_bit, + expert_parallel_group=expert_parallel_group, + moe_capacity_factor=moe_capacity_factor, + moe_expert_model_parallelism=has_parallel, + forward_fn=forward_fn, + hidden_size=self.experts.hidden_size, + mlp_impl=mlp_impl, + # Shared expert parameters + shared_up_proj_weight=self.shared_up_proj_weight, + shared_down_proj_weight=self.shared_down_proj_weight, + shared_up_proj_bias=self.shared_up_proj_bias, + shared_down_proj_bias=self.shared_down_proj_bias, + shared_expert_weighted_sum=self.shared_expert_weighted_sum, + shared_activation_fn=self.shared_activation_fn, + ) + return output, expert_weights_out + + +# Patch for XPU or CPU support +if hasattr(torch, "xpu") and torch.xpu.is_available(): + from .xpu_fused_moe import MegaBlocksMoeMLP + +from .cpu_moe_cpp import CPUMegaBlocksMoeMLP diff --git a/build/torch214-cxx11-cu130-x86_64-linux/metadata.json b/build/torch214-cxx11-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..7608c8a414755734e4a48d1a503b0417f22c223c --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/metadata.json @@ -0,0 +1,120 @@ +{ + "name": "megablocks", + "id": "_megablocks_cuda_addf474", + "version": 1, + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.3", + "11.0", + "12.0", + "12.1+PTX", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "o7Y9XLXezcJogNTNoC4HMHsd2ZZ3CTwzYftkUthOtno=", + "_grouped_gemm_triton/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "_grouped_gemm_triton/_triton_kernels/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "_grouped_gemm_triton/_triton_kernels/gmm.py": "11SvNqMt/LDCgjJBOmZlHg9yn9NbGw1UYKzNR/JKtcQ=", + "_grouped_gemm_triton/adapter.py": "tP00AypBJn4rDqrBDMMwBb6MwJ6s9dg18GNQxpLiVoA=", + "_grouped_gemm_triton/configs.py": "aHXlKELHIy8J82m1kClQRxDkad4LCL8BOeEpq4/IqKQ=", + "_grouped_gemm_triton/gmm.py": "1iqK7N3vsahgQ4bbHsOMA8hoR5uQhh1PmbFWjOLSKDo=", + "_grouped_gemm_triton/utils/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "_grouped_gemm_triton/utils/_triton/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "_grouped_gemm_triton/utils/_triton/arch_info.py": "Vo7tziqn9GLA8S4xrKT2555aKFxo75Oja4ek2zNReuY=", + "_grouped_gemm_triton/utils/_triton/pid_preprocessing.py": "Lx1S7li2xArbeql/pEFfyul7yU61agwuYHTsZ7jG0os=", + "_grouped_gemm_triton/utils/gmm_common.py": "pQMeVpP9kk/bixwJE309iceJPa8H7mZV8b6+OtbQvl8=", + "_grouped_gemm_triton/utils/logger.py": "CXuJZCBzQDTrW80OyJ2IuhbHhayLy12Ake1KGxiuYIA=", + "_layers/__init__.py": "bwEDLk8kbLuT1BWZPPDVUiDsWsJ3wC3nI+O+dSGAoSU=", + "_layers/activation_fn.py": "lJrZcw+1oLw0XOyjYb6AfmxX1n6g2Snbph6H3p5D0oU=", + "_layers/all_to_all.py": "SZMHtoTy6B76caEE0shoK1y4iHqWp+lWqdR/30d9oBQ=", + "_layers/arguments.py": "wYFjDwBidH5MjMr4cZGDEcRvW1IgdLQvcYt/qt5S0dk=", + "_layers/common.py": "8m+bNabk+jB+C8o797GdNpGzx8m3nbnG66FQdnDZRVI=", + "_layers/dmlp_registry.py": "z7iqmxiJx9eACdaHOCmnY9EGKV/zePhqcf8FVboaoK8=", + "_layers/dmoe.py": "bOaQ0PFOOLu45HxJNRFuKPB8VedUzsduekkqJV6WrQ8=", + "_layers/gelu.py": "fcq4OBCSCcPjCk7PYZydIpbdn/iUyMtL2TqhonHqt7k=", + "_layers/glu.py": "D27/lTdF05N6TdSPoyW1MWgTGyaXJlUCo5uos54zabU=", + "_layers/memory_test.py": "4jm72yP6WBV5zelWQGJQIHHZLPisFEtar+moOSnxXMk=", + "_layers/mlp.py": "xx51A1pGeBt/B8m7pI5hjhVqhEllpxvldulXYzFmA3Y=", + "_layers/moe.py": "vVqEx4Ppgw+QbLHOJo5uSI/qXDikA99ZGApPOcBJtmg=", + "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=", + "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=", + "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=", + "_megablocks_cuda_addf474.abi3.so": "q3FYXoFtsZaVjPbaqvv+Ve1JxR5OvJuUmMALQdWD9p0=", + "_ops.py": "WXlV4XwtKQjdQ6nap3k9TalaHkK0ypeSTshOrV6Hczo=", + "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=", + "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=", + "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=", + "benchmark_util.py": "WFsGWGaJvWvwdx4w33WCfNzSJkVuQrZmZCaAodJsws8=", + "cpu_fused_moe.py": "XN2eQBFaihbJrZ/cu9jS8h5dcoiUNh82gjJQUARcBYA=", + "cpu_moe_cpp.py": "dvv4Ynn/FWObs421L4hAdYZPiPs2BS/sbEOuMMIUPVM=", + "grouped_gemm/__init__.py": "fxW9amtzy/Z0V/OXSYMmibD7XvpYBGEv81KMxvocFKY=", + "grouped_gemm/backend.py": "Yd//ihpvu/pxEoGBq3iD1zMggnYv6quwZMdg1R4kUEI=", + "grouped_gemm/ops.py": "Kj7yRgv7afB5SVwXIkN9RNIF2kxD2qrZ0dVsKX5ph2s=", + "grouped_gemm_util.py": "bVPrdsEtBVJJ+v/bNrGNf/QHNlVL4VxHg92HvaL3oq4=", + "layers.py": "mNTlX2oKJyHBJllRd24lPNdGLa2H9cEUgVLvkZJpoNo=", + "ops/__init__.py": "MHQsiMEgkKeZCa8TAYFMGDFnwFIgijfAN7DhMUt/48U=", + "ops/all_to_all_benchmark.py": "sxXSCSDdsh+62IWpKNTyBphLTngY1jlgJ1WgNooSEO8=", + "ops/binned_gather.py": "pkWne+ThfiTR/bMbDs2QDimKxAI4klqD1FdvhytE7lU=", + "ops/binned_scatter.py": "5NxocRmgYDkTAntDtXBlF4c2r8rfbLVnBRPA6f4/L4U=", + "ops/cumsum.py": "upuWnnp125kQhJrIMLtG09cI1zzxEy06P5wYhOwbpI4=", + "ops/gather.py": "a8j/btP1VCdJ7F7FC58MFPxFuH9tUhjMuQnEzwjAgGw=", + "ops/histogram.py": "v5+pHF+/fjKQAkNLIsHyljMUakMTdW5GzA094PgRoGk=", + "ops/histogram_benchmark.py": "oX8VolN+/PVCWaPHkY0k7SNn4QpadHi80Sla/3NScNI=", + "ops/matmul_benchmark.py": "SZ3L/CatnT4J5PyOU/Ii/b5y5fw3HNEaR/xY6PjrJ8A=", + "ops/padded_gather.py": "VtPH3Ge1jodeug5MWwXnxnmGIi0D1MOOaXi2CcU8lpI=", + "ops/padded_scatter.py": "UO1BwsaezK3AikD8CHoJOEjmC4h+H5fcOX/bvO8FWXw=", + "ops/padded_scatter_benchmark.py": "w34qlIc3Nn0AgB2k7iIFqQXgDNplMVqCr5i8GrgrvZ0=", + "ops/permute_benchmark.py": "ISYPkVZknTnfB8fhUVloSyfEOe7ApKQKtNPvEq3P9qg=", + "ops/repeat.py": "j9l1A/4LVhcFuSAaneFlImq7RGNXie0zfu2nMlEhHI4=", + "ops/replicate.py": "NN5843jJz0OdfYS6YaPCpRD8dGk66S3SHTzR51PrwUI=", + "ops/round_up.py": "/CQSCYmjW0Mrv8TFdfZs/ZpXHHz18b3E4aLHvmfUxOY=", + "ops/scatter.py": "YgM76I/dYAzRNN+jzAdaRWan56GiwuCZRR7xy3RAY8Y=", + "ops/sort.py": "35D3oG4AT0Jd6+awwfoL+GmkcNtdRLno7LcSFidqKts=", + "ops/sort_benchmark.py": "Bo985Xsxq/F0PMU92xKCo+srSKVdheQLey2GqFlwZXY=", + "ops/stk_autocast.py": "SQArC9BFDXVManccbEWibSg0E+/R3NQlkp0cB4mIjXo=", + "ops/sum.py": "fcjWM5Aiz1wQ+UYBN2Y/OFILsiGKSLgmlkeqfVRccD4=", + "ops/topology.py": "+lfjfxDjB+st3K8wZnRY5fEanSIWT+3SJncx9bu+6P0=", + "stk/__init__.py": "QFugfJQ5q1Wdn/+VQyfs1YKns79i3FiywiLevg0ymXo=", + "stk/backend/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "stk/backend/autocast.py": "ppk7mL/WQUHaLqEgsW3h5ZGhldC8WyIKmaIS6kAe4lA=", + "stk/backend/sputnik.py": "Mpbo0M1ES39PdaIo65ef3/mQ8XHtLbO3ayqufaIZLiw=", + "stk/backend/triton_kernels.py": "RCthwX2feppkqBK3pvm+/37nhF89lLHgP12c6xrfxXU=", + "stk/matrix.py": "Xxo1RwOtKFgVXExhlAoIC67V+eWKZRClDdbBfzUgfy0=", + "stk/ops/__init__.py": "Vqbt6C0qWEPaWypsCoaHGoZRbdPcFrLg1ahlAlp4l+Y=", + "stk/ops/eltwise_ops.py": "O9TqOMvZ2NZ7KNFaeuCmJOfg1Hz5ad8rvYwbGAH7/+0=", + "stk/ops/eltwise_ops_test.py": "rD1oPPu+vXHlyJnEJOpM9L5wdCeiDbn98CbCH8mQ0wE=", + "stk/ops/linear_ops.py": "ZWOpLTXMdw/ZTyxL4W74CSSZQaUKvtfiz3ev0OnaVQ0=", + "stk/ops/linear_ops_test.py": "kl8nMX1EAvZ6lF7oN4U4bW2z0hera0YxqTPiNljLmzk=", + "stk/ops/matrix_ops.py": "+hhUk7IZZFVVzIk1pttWFFSOtKwlGEBVR2V4eyL27K4=", + "stk/ops/matrix_ops_test.py": "Bb4fk4+IzNC0zOFq0Tz/JwNaLWyhtWEDgEUzyVdbnrE=", + "stk/random/__init__.py": "Gr/+Gl20jJ6fJVRfoAUxEcCPeChILQzrzhKVVAWQsRI=", + "stk/random/random_ops.py": "AluIDj3Vo2ReKAsZz07yydbtndDbJoglekRGIaCoRZE=", + "stk/random/random_ops_test.py": "B+a4Zb+OPhBtpVA48oc09sBL/AAQO9j/YUFI0nnXxzk=", + "xpu_fused_moe.py": "sPmuRSpkH1Bmk4hdpz53us8Qpc6yinT9rkld3T8WreI=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "commit": "7d2828e6592de641b44fb8eb896719101a2ab2fb", + "dirty": false + }, + "kernel": { + "commit": "addf4741fafb55b4c101b15985dfaeea7c7c2f3a", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch214-cxx11-cu130-x86_64-linux/metadata.json.sigstore b/build/torch214-cxx11-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..ab90fea4412d4be94b7251e01b0227c5396a77b4 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHTDCCBtGgAwIBAgIUKFrcoV0SCAIo4SZry67kyp0v1zQwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI3MDcxNDA4WhcNMjYwODI3MDcyNDA4WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEZQ7EY+AW7pTYSKJOvoFRFPKSrxCylI1CXd3MCOi/ZvYkUbXqCir8pF59+VDe0FljxqQzkbRH7RICGNl7+RBwMqOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUWSKtOGUYBitWLuC83MjG1XeEfSswHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGFkZGY0NzQxZmFmYjU1YjRjMTAxYjE1OTg1ZGZhZWVhN2M3YzJmM2EwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzMwNDUzOTU5NTIvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoEIRqGEAAAQDAEcwRQIgRz5E4gbFKvIi2Q47D/XHOEqUTdf88vAW/DtYYyMKrbQCIQCP7D4x+thTWsp7r/1bCQDXp4+Jrannj9VxIYqjGUvZUDAKBggqhkjOPQQDAwNpADBmAjEAjxJ1dLH53pH2gAKKx/PuEVOPapt0lJabWJd75AWdGa40/FzAo5++SJL8s+GAMCqwAjEApAMiuD6veMAcOCNWR/gFLhHj4src6SIdVeojFiIBmxJP4CwPIRDrPq2tMKKql6AH"}, "tlogEntries":[{"logIndex":"2614047672", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787814848", "inclusionPromise":{"signedEntryTimestamp":"MEUCIGy1VZsg8tkERkm91xtb0b+97cFEq58PsLOpP3TRAn+1AiEAtnx/47vSsBKnKtYDWiVPGYOVGKh0mFxnkCw+GWPN56o="}, "inclusionProof":{"logIndex":"2492143410", "rootHash":"YT4XDTCCBqJK6grnyMSn6+6Mwo1C2f4wwqFZ8HplHWM=", "treeSize":"2492143412", "hashes":["x+P1PG8t9U/ngjOQAScoaQvsrbjO+bGQx/bxIhaP+EM=", "emONeWjYzZSGQr5FKjRPtsFmxgVa/zUBmUsZnJ/JoJU=", "S5RKLg6QeqLGcjqNOO7Kh6ksGe2E9/w65+VyfNmtHTU=", "zaAwsvMTV+3TE24aN0KDcD60+iJAlqM6VMij7WK/CPo=", "6oMuxhmD5eaznFmW0FxN8Kz33TK8kJ0t4lDTtjU4HIs=", "74YOSoxcaARV3QmRlp2HcIEFoCumNC8/rDR31M3b3Zk=", "KIag3bHIhHqGqEJl+No8RWATx3rcwFMD/MdDkBYHPmQ=", "CFNMzmJ8N+uNx37P2asMqvrW9tw4axAVHnXLRO39tFs=", "NypqAbB+mNix4l1ahGeP9fRWiQzZCm4P1664makjjq4=", "8SoaeA7IeJPoIiveQbBfHCWXSDQabIk2PW+FT49NSEQ=", "VGrFt5JvlIUQvoGEdBBxrCe3lYbu581OhgLq/Xdyat4=", "7TfDCaMWgEwSJEZd/e33b24919rF4IDJLmbqNNRNhZs=", "0HfdxDt/zGugZLIuHasrdLEW9s8OaTykSSqBEDFQXdw=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2492143412\nYT4XDTCCBqJK6grnyMSn6+6Mwo1C2f4wwqFZ8HplHWM=\n\n— rekor.sigstore.dev wNI9ajBEAiB9LJZPA1K8FQXyw+7rlaEoX2bzAEw/eNaZHXJY18+kmQIgDwVYTRrPxB5EnnG6Km+gqxzwd4TH0RY5Yx95NWsKRc8=\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI4YWVmMjk2YmQ0YjU2YzZiNzZiN2RmMjQ0ZmQ1NjZkZDdjNzkzZTY0NzRhMzg5MWY1ZDllYTg5YmI5MjI2YzAwIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRnBVdHA1ZDA5bXRWNFV1dFh4OWtnYVFMMmFYYko5QkFhUlB5SzZySFRKSUFpRUFvdUcxZUJpcUNsVVhBR2s4eEhONmFqQVRSN2ZqTzdtREx1d2QrcDRsU01nPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVSRU5EUW5SSFowRjNTVUpCWjBsVlMwWnlZMjlXTUZORFFVbHZORk5hY25rMk4ydDVjREIyTVhwUmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1ROTlJHTjRUa1JCTkZkb1kwNU5hbGwzVDBSSk0wMUVZM2xPUkVFMFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZhVVRkRldTdEJWemR3VkZsVFMwcFBkbTlHVWtaUVMxTnllRU41YkVreFExaGtNMDBLUTA5cEwxcDJXV3RWWWxoeFEybHlPSEJHTlRrclZrUmxNRVpzYW5oeFVYcHJZbEpJTjFKSlEwZE9iRGNyVWtKM1RYRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZYVTB0MENrOUhWVmxDYVhSWFRIVkRPRE5OYWtjeFdHVkZabE56ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVpWMUpyV21wUk0wNUVSbTFaVjFwcFRsUldhVTVIVFhoTlJFWnBDazFVVlRWUFJGWnJXbTFHYkZwWFJUTlplbVJxVFcxWmVsbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWxOVkZVMVQwUldhMXB0Um13S1dsZEZNMWw2WkdwTmJWbDZXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWtLVFZSVk5VOUVWbXRhYlVac1dsZEZNMWw2WkdwTmJWbDZXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZEdhMXBIV1RBS1RucFJlRnB0Um0xWmFsVXhXV3BTYWsxVVFYaFpha1V4VDFSbk1WcEhXbWhhVjFab1RqSk5NMWw2U20xTk1rVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWsxM1RrUlZlazlVVlRWT1ZFbDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUZTVkp4UjBWQlFVRlJSQXBCUldOM1VsRkpaMUo2TlVVMFoySkdTM1pKYVRKUk5EZEVMMWhJVDBWeFZWUmtaamc0ZGtGWEwwUjBXVmw1VFV0eVlsRkRTVkZEVURkRU5IZ3JkR2hVQ2xkemNEZHlMekZpUTFGRVdIQTBLMHB5WVc1dWFqbFdlRWxaY1dwSFZYWmFWVVJCUzBKblozRm9hMnBQVUZGUlJFRjNUbkJCUkVKdFFXcEZRV3A0U2pFS1pFeElOVE53U0RKblFVdExlQzlRZFVWV1QxQmhjSFF3YkVwaFlsZEtaRGMxUVZka1IyRTBNQzlHZWtGdk5Tc3JVMHBNT0hNclIwRk5RM0YzUVdwRlFRcHdRVTFwZFVRMmRtVk5RV05QUTA1WFVpOW5Sa3hvU0dvMGMzSmpObE5KWkZabGIycEdhVWxDYlhoS1VEUkRkMUJKVWtSeVVIRXlkRTFMUzNGc05rRklDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgf2Y7wyq0VEO/koaXbQGiMZVihHodM4HbLuveG0n/u9ECFDqkTyZaoiGTkH0Y7auV1PJeDNcKGA8yMDI2MDgyNzA3MTQwOFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwODI3MDcxNDA4WjAvBgkqhkiG9w0BCQQxIgQg3crqx2rrbSUFI4Wm/1Ib7I4RxKEYlhO5WIZZ0ZBWfiYwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIwW9ZCGYM8iCcyBkSfDMcFyK0WgUrwJ9YeA7MWP+bQ1PCnCZcj9qdxEmCoDP/b4gDyAjEA4cEh2jKOu0XBrs80QYXV9dstscI5AJoBOWO0PN7MM3Nu+r7lftLIUaGasAyRX2tM"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"iu8pa9S1bGt2t98kT9Vm3Xx5PmR0o4kfXZ6om7kibAA="}, "signature":"MEUCIFpUtp5d09mtV4UutXx9kgaQL2aXbJ9BAaRPyK6rHTJIAiEAouG1eBiqClUXAGk8xHN6ajATR7fjO7mDLuwd+p4lSMg="}} \ No newline at end of file diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b944080df810d0b0cfc571f3009b0098a651f9b7 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/__init__.py @@ -0,0 +1,35 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from .binned_gather import binned_gather +from .binned_scatter import binned_scatter +from .cumsum import exclusive_cumsum, inclusive_cumsum +from .gather import gather +from .histogram import histogram +from .padded_gather import padded_gather +from .padded_scatter import padded_scatter +from .repeat import repeat +from .replicate import replicate +from .round_up import round_up +from .scatter import scatter +from .sort import sort +from .sum import sum +from .topology import topology + +__all__ = [ + 'binned_gather', + 'binned_scatter', + 'exclusive_cumsum', + 'inclusive_cumsum', + 'gather', + 'histogram', + 'padded_gather', + 'padded_scatter', + 'repeat', + 'replicate', + 'round_up', + 'scatter', + 'sort', + 'sum', + 'topology', +] diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/all_to_all_benchmark.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/all_to_all_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..4c939818edca3345f6344bbc7cef07ffe3cd0181 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/all_to_all_benchmark.py @@ -0,0 +1,63 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.distributed as dist + +# from megablocks import benchmark_util +# from megablocks.layers.all_to_all import all_to_all + +from .. import benchmark_util +from .._layers.all_to_all import all_to_all + +_ALL_TO_ALL_BENCHMARK = ( + (8, 1024), + (16, 1024), + (32, 1024), + (64, 1024), + (128, 1024), + (256, 1024), + (512, 1024), + (1024, 1024), + (2 * 1024, 1024), + (4 * 1024, 1024), + (8 * 1024, 1024), + (16 * 1024, 1024), + (32 * 1024, 1024), + (64 * 1024, 1024), + (128 * 1024, 1024), + (256 * 1024, 1024), + (512 * 1024, 1024), + (1024 * 1024, 1024), +) + + +def benchmark_all_to_all(group, sl, hs): + world_size = dist.get_world_size(group) + assert (sl % world_size) == 0 + send_recv_sizes = [sl // world_size] * world_size + + x = torch.randn((sl, hs)).cuda().half() + + details = { + 'world_size': world_size, + 'message_size (B)': send_recv_sizes[0] * hs * 2, # 2B elements. + } + + def benchmark(): + return all_to_all(x, send_recv_sizes, send_recv_sizes, group) + + time, std = benchmark_util.benchmark_function(benchmark) + + if dist.get_rank(group) == 0: + benchmark_util.log_benchmark('All-To-All', details, time, std) + + +if __name__ == '__main__': + assert dist.is_available() + group = dist.init_process_group(backend='nccl') + local_rank = dist.get_rank(group) + torch.cuda.set_device(local_rank) + + for args in _ALL_TO_ALL_BENCHMARK: + benchmark_all_to_all(group, *args) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/binned_gather.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/binned_gather.py new file mode 100644 index 0000000000000000000000000000000000000000..189a7fa3518d660f29ea32e7a04827164af98d60 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/binned_gather.py @@ -0,0 +1,37 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for binned_gather kernel. +class BinnedGatherOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bins: torch.Tensor, + bin_size: int, + top_k: int, + ): + ctx.save_for_backward(indices, bins) + ctx.top_k = top_k + return kernels.binned_gather(x, indices, None, bins, bin_size, top_k) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + indices, bins = ctx.saved_tensors + out = kernels.binned_scatter(grad, indices, None, bins, ctx.top_k) + return out, None, None, None, None + + +binned_gather = BinnedGatherOp.apply diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/binned_scatter.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/binned_scatter.py new file mode 100644 index 0000000000000000000000000000000000000000..cb937c0c106662ce8108c1cb926f8f063b163d3d --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/binned_scatter.py @@ -0,0 +1,59 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for binned_scatter kernel. +class BinnedScatterOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + top_k: int, + ): + assert len(x.size()) == 3 + ctx.bin_size = x.size(1) + ctx.top_k = top_k + + # TODO(tgale): Don't save 'x' for backwards if we don't need to + # calculate the gradient w.r.t. 'weights'. + ctx.save_for_backward(x, indices, weights, bins) + return kernels.binned_scatter(x, indices, weights, bins, top_k) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + x, indices, weights, bins = ctx.saved_tensors + out = kernels.binned_gather( + grad, + indices, + weights, + bins, + ctx.bin_size, + ctx.top_k, + ) + + wgrad = None + if ctx.needs_input_grad[2]: + wgrad = kernels.binned_scatter_wgrad( + x, + grad, + indices, + bins, + ctx.top_k, + ) + return out, None, wgrad, None, None + + +binned_scatter = BinnedScatterOp.apply diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/cumsum.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/cumsum.py new file mode 100644 index 0000000000000000000000000000000000000000..e2b7572391e20045d335cf7337246e8a9b9f57ef --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/cumsum.py @@ -0,0 +1,52 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + # import megablocks_ops as ops # type: ignore + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + + +# Autograd wrappers for cumsum kernels. +# NOTE: Does not support gradients. +class ExclusiveCumsumOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, dim: int): + if len(x.size()) == 1: + x = x.view([1, -1]) + out = torch.empty_like(x) + ops.exclusive_cumsum(x, 1, out) + return out.squeeze() + out = torch.empty_like(x) + ops.exclusive_cumsum(x, dim, out) + return out + + +exclusive_cumsum = ExclusiveCumsumOp.apply + + +class InclusiveCumsumOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, dim: int) -> torch.Tensor: + if len(x.size()) == 1: + x = x.view([1, -1]) + out = torch.empty_like(x) + ops.inclusive_cumsum(x, 1, out) + return out.squeeze() + out = torch.empty_like(x) + ops.inclusive_cumsum(x, dim, out) + return out + + +inclusive_cumsum = InclusiveCumsumOp.apply diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/gather.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/gather.py new file mode 100644 index 0000000000000000000000000000000000000000..f1f87c1e7bed8d3589dd790805234976e0b05898 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/gather.py @@ -0,0 +1,38 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for gather kernel. +class GatherOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + bins: torch.Tensor, + top_k: int, + ): + ctx.save_for_backward(indices, bin_ids, bins) + ctx.top_k = top_k + return kernels.gather(x, indices, bin_ids, None, bins, top_k) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + + indices, bin_ids, bins = ctx.saved_tensors + out = kernels.scatter(grad, indices, bin_ids, None, bins, ctx.top_k) + return out, None, None, None, None, None + + +gather = GatherOp.apply diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/histogram.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/histogram.py new file mode 100644 index 0000000000000000000000000000000000000000..7b3f058ec373cbba7555704fb5e4212c3cc75d9d --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/histogram.py @@ -0,0 +1,27 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + + +# Autograd wrapper for histogram kernel. +# NOTE: Does not support gradients. +class HistogramOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, max_val: float): + return ops.histogram(x, max_val) + + +histogram = HistogramOp.apply diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/histogram_benchmark.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/histogram_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..a0b9c6047567b87a295979498142230d1b0c9411 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/histogram_benchmark.py @@ -0,0 +1,78 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import numpy as np +import torch +# from absl.testing import parameterized + +from .. import ops + +_HISTOGRAM_TESTS = ( + (16384, torch.int32, 2), + (16384, torch.int32, 4), + (16384, torch.int32, 8), + (16384, torch.int32, 16), + (16384, torch.int32, 32), + (16384, torch.int32, 64), + (16384, torch.int32, 128), + (16384, torch.int32, 256), +) + + +def benchmark_function(fn, iterations=10): + # Run once to get rid of startup overhead. + fn() + times = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + times = np.array(times) + return times.mean(), times.std(), times.max(), times.min() + + +def log_benchmark(arguments, mean_t, std_t): + print('=' * 60) + print('Benchmark Parameters:') + for (key, value) in arguments.items(): + print(f'{key} = {value}') + print('Results:') + print('mean / std = {:.2f}ms / {:.2f}ms'.format(mean_t, std_t)) + print('=' * 60) + + +# class HistogramBenchmark(parameterized.TestCase): +# +# @parameterized.parameters(*_HISTOGRAM_TESTS) +# def testHistogram(self, n, dtype, max_val): +# x = torch.randint(0, max_val, (n,)).cuda().to(dtype) +# +# mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.histogram(x, max_val),) +# arguments = { +# 'n': n, +# 'dtype': dtype, +# 'max_val': max_val, +# } +# log_benchmark(arguments, mean_t, std_t) +# +# @parameterized.parameters(*_HISTOGRAM_TESTS) +# def testTorchHistogram(self, n, dtype, max_val): +# x = torch.randint(0, 128, (n,)).cuda().to(dtype) +# +# mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.histc(x, max_val, 0, max_val - 1),) +# arguments = { +# 'n': n, +# 'dtype': dtype, +# 'max_val': max_val, +# } +# log_benchmark(arguments, mean_t, std_t) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/matmul_benchmark.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/matmul_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..4c4b9b8866ffed2eb769b77f2320c82e5491ae0e --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/matmul_benchmark.py @@ -0,0 +1,415 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + + +# import stk + +# try: +# import stk +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/ops/matmul_benchmark.py is needed.', +# ) + +from .. import stk + +import torch +# from absl.testing import parameterized + +from .. import benchmark_util, ops + + +# Calling tensor.t() calls tensor.transpose(0, 1) which calls +# torch.as_strided(...). Circumvent this chain to avoid an overhead +# this adds. +def transpose_view(x): + return torch.as_strided( + x, + (x.shape[1], x.shape[0]), + (x.stride()[1], x.stride()[0]), + ) + + +_MATMUL_TESTS = ( + (64 * 1024, 512, 2048, 64), + (32 * 1024, 768, 3072, 64), + (8 * 1024, 1024, 4096, 64), + (4 * 2048, 4096, 4 * 4096, 4), +) + + +def log_benchmark(name, arguments, time, std, flops): + benchmark_util.log_benchmark(name, arguments, time, std) + print('flops = {:.2f}B'.format(flops / 1e9)) + print('throughput = {:.2f}T'.format(flops / 1e9 / time)) + print('=' * 60) + + +# class MatmulBenchmark(parameterized.TestCase): +# +# def build_sparse_matrix(self, x, padded_bins, fhs, ne): +# blocking = 128 +# padded_tokens, _ = x.size() +# assert padded_tokens % blocking == 0 +# assert fhs % blocking == 0 +# +# # Offsets for the sparse matrix. All rows have the +# # same number of nonzero blocks dictated by the +# # dimensionality of a single expert. +# block_rows = padded_tokens // blocking +# blocks_per_row = fhs // blocking +# offsets = torch.arange( +# 0, +# block_rows * blocks_per_row + 1, +# blocks_per_row, +# dtype=torch.int32, +# device=x.device, +# ) +# +# # Indices for the sparse matrix. The indices for +# # the intermediate matrix are dynamic depending +# # on the mapping of tokens to experts. +# column_indices = ops.topology( +# padded_bins, +# blocking, +# block_rows, +# blocks_per_row, +# ) +# data = torch.empty( +# column_indices.numel(), +# blocking, +# blocking, +# dtype=torch.float16, +# device=x.device, +# ) +# shape = (padded_tokens, fhs * ne) +# row_indices = stk.ops.row_indices(shape, data, offsets, column_indices) +# return stk.Matrix(shape, data, row_indices, column_indices, offsets) +# +# def build_input_matrix(self, sl, hs, ne): +# x = torch.randn((sl, hs)).cuda().half() +# +# # Assign tokens to experts uniformly. +# top_expert = torch.arange(0, sl).cuda().int() % ne +# +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(top_expert, ne) +# padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) +# padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# out = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, 1) +# return out, padded_bins +# +# def build_weight_matrix(self, ne, hs, fhs): +# return torch.randn((hs, ne * fhs)).cuda().half() +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_Fwd_SDD_NT(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# topo = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# w = transpose_view(w) +# +# def benchmark(): +# return stk.ops.sdd(x, w, topo) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0::Fwd::SDD::NT', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_GradX_DSD_NN(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# topo = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# +# def benchmark(): +# return stk.ops.dsd(topo, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0::GradX::DSD::NN', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_GradW_DSD_TN(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# topo = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# topo = topo.t() +# +# def benchmark(): +# return stk.ops.dsd(topo, x) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0::GradW::DSD::TN', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_Fwd_DSD_NN(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# x = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# +# def benchmark(): +# return stk.ops.dsd(x, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::Fwd::DSD::NN', +# arguments, +# mean_t, +# std_t, +# x.nnz * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_GradX_SDD_NT(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# x = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# out = stk.ops.dsd(x, w) +# w = transpose_view(w) +# +# def benchmark(): +# return stk.ops.sdd(out, w, x) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::GradX::SDD::NT', +# arguments, +# mean_t, +# std_t, +# x.nnz * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_GradW_DSD_TN(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# x = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# out = stk.ops.dsd(x, w) +# x = x.t() +# +# def benchmark(): +# return stk.ops.dsd(x, out) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::GradW::DSD::TN', +# arguments, +# mean_t, +# std_t, +# x.nnz * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_Fwd_DDD_NT(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, hs)).cuda().half() +# w = torch.randn((ne, hs, fhs)).cuda().half() +# +# w = w.transpose(1, 2).contiguous() +# w = w.transpose(1, 2) +# +# def benchmark(): +# return torch.bmm(x, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0::Fwd:DDD::NT', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_GradX_DDD_NN(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, hs)).cuda().half() +# w = torch.randn((ne, hs, fhs)).cuda().half() +# out = torch.bmm(x, w) +# w = w.transpose(1, 2).contiguous() +# +# def benchmark(): +# return torch.bmm(out, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0:GradX:DDD::NN', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_GradW_DDD_TN(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, hs)).cuda().half() +# w = torch.randn((ne, hs, fhs)).cuda().half() +# out = torch.bmm(x, w) +# out = out.transpose(1, 2) +# +# def benchmark(): +# return torch.bmm(out, x) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0:GradW:DDD::TN', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_Fwd_DDD_NN(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, fhs)).cuda().half() +# w = torch.randn((ne, fhs, hs)).cuda().half() +# +# def benchmark(): +# return torch.bmm(x, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::Fwd::DDD::NN', +# arguments, +# mean_t, +# std_t, +# x.numel() * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_GradX_DDD_NT(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, fhs)).cuda().half() +# w = torch.randn((ne, fhs, hs)).cuda().half() +# out = torch.bmm(x, w) +# w = torch.transpose(w, 1, 2) +# +# def benchmark(): +# return torch.bmm(out, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::GradX::DDD::NT', +# arguments, +# mean_t, +# std_t, +# x.numel() * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_GradW_DDD_TN(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, fhs)).cuda().half() +# w = torch.randn((ne, fhs, hs)).cuda().half() +# out = torch.bmm(x, w) +# x = torch.transpose(x, 1, 2) +# +# def benchmark(): +# return torch.bmm(x, out) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::GradW::DDD::TN', +# arguments, +# mean_t, +# std_t, +# x.numel() * hs * 2, +# ) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/padded_gather.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/padded_gather.py new file mode 100644 index 0000000000000000000000000000000000000000..c1cf4047c9494394d2a3884ba8830179013db7ff --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/padded_gather.py @@ -0,0 +1,55 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for padded_gather kernel. +class PaddedGatherOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + bins: torch.Tensor, + padded_bins: torch.Tensor, + top_k: int, + ): + ctx.save_for_backward(indices, bin_ids, bins, padded_bins) + ctx.top_k = top_k + return kernels.padded_gather( + x, + indices, + bin_ids, + None, + bins, + padded_bins, + top_k, + ) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + + indices, bin_ids, bins, padded_bins = ctx.saved_tensors + out = kernels.padded_scatter( + grad, + indices, + bin_ids, + None, + bins, + padded_bins, + ctx.top_k, + ) + return out, None, None, None, None, None + + +padded_gather = PaddedGatherOp.apply diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/padded_scatter.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/padded_scatter.py new file mode 100644 index 0000000000000000000000000000000000000000..61e021b81497e472cda5d72bdac557a0ca92d262 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/padded_scatter.py @@ -0,0 +1,98 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for padded_scatter kernel. +class PaddedScatterOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + padded_bins: torch.Tensor, + top_k: int, + ): + maybe_x = [x] if ctx.needs_input_grad[3] else [] + ctx.save_for_backward( + indices, + bin_ids, + weights, + bins, + padded_bins, + *maybe_x, + ) + ctx.top_k = top_k + ctx.x_shape = x.shape + return kernels.padded_scatter( + x, + indices, + bin_ids, + weights, + bins, + padded_bins, + top_k, + ) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + saved_tensors = ctx.saved_tensors + + indices, bin_ids, weights, bins, padded_bins = saved_tensors[:5] + dgrad = None + if ctx.needs_input_grad[0]: + dgrad = kernels.padded_gather( + grad, + indices, + bin_ids, + weights, + bins, + padded_bins, + ctx.top_k, + ) + + wgrad = None + if ctx.needs_input_grad[3]: # need wgrad + x = saved_tensors[-1] + wgrad = kernels.padded_scatter_wgrad( + x, + grad, + indices, + bin_ids, + bins, + padded_bins, + ctx.top_k, + ) + return dgrad, None, None, wgrad, None, None, None, None + + +def padded_scatter( + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + padded_bins: torch.Tensor, + top_k: int, +): + return PaddedScatterOp.apply( + x, + indices, + bin_ids, + weights, + bins, + padded_bins, + top_k, + ) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/padded_scatter_benchmark.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/padded_scatter_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..3cbe4735891446b46f93170c64c23fe63632bf93 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/padded_scatter_benchmark.py @@ -0,0 +1,66 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import torch +# from absl.testing import parameterized + +from .. import benchmark_util, ops + +_PADDED_SCATTER_BENCHMARK = ( + # dMoE-Medium, 8-way EMP. + (1024 * 16, 1024, 8, 4), + # dMoE-Medium, post-all-to-all. + (1024 * 16 * 4, 1024, 8, 1), +) + + +# class PaddedScatterTest(parameterized.TestCase): +# +# @parameterized.parameters(*_PADDED_SCATTER_BENCHMARK) +# def testPaddedScatter(self, sl, hs, ne, top_k): +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# +# # Randomly assign tokens to experts. +# top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(top_expert, ne) +# padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) +# padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# +# # Sample weights for the scatter reduce. +# weights = torch.rand((sl * top_k,)).cuda().half() +# +# # Gather the data to prepare for backwards. +# x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k) +# +# def benchmark(): +# return ops.padded_scatter( +# x, +# indices, +# bin_ids, +# weights, +# bins, +# padded_bins, +# top_k, +# ) +# +# time, std = benchmark_util.benchmark_function(benchmark) +# benchmark_util.log_benchmark( +# 'Padded Scatter', +# { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# 'top_k': top_k, +# }, +# time, +# std, +# ) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/permute_benchmark.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/permute_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..697abddbb3a2082ec4ddd6d94f89f7faabb34b40 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/permute_benchmark.py @@ -0,0 +1,149 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import torch +# from absl.testing import parameterized + +from .. import benchmark_util, ops + +_PERMUTE_TESTS = ( + (16384, 768, 2), + (16384, 768, 4), + (16384, 768, 8), + (16384, 768, 16), + (16384, 768, 32), + (16384, 768, 64), + (16384, 768, 128), + (16384 * 8, 768, 2), + (16384 * 8, 768, 4), + (16384 * 8, 768, 8), + (16384 * 8, 768, 16), + (16384 * 8, 768, 32), + (16384 * 8, 768, 64), + (16384 * 8, 768, 128), +) + + +# class PermuteBenchmark(parameterized.TestCase): +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testBinnedGather(self, sl, hs, ne): +# # NOTE: Capacity factor == 1. +# ec = sl // ne +# +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# top_expert = torch.randint(0, ne, (sl,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(indices, ne) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# +# def benchmark(): +# return ops.binned_gather(x, indices, bins, ec) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('BinnedGather', arguments, mean_t, std_t) +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testBinnedScatter(self, sl, hs, ne): +# # NOTE: Capacity factor == 1. +# ec = sl // ne +# +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# top_expert = torch.randint(0, ne, (sl,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(indices, ne) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# x = ops.binned_gather(x, indices, bins, ec) +# +# def benchmark(): +# return ops.binned_scatter(x, indices, bins) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('BinnedScatter', arguments, mean_t, std_t) +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testPaddedGather(self, sl, hs, ne): +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# +# # Randomly assign tokens to experts. +# top_expert = torch.randint(0, ne, (sl,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(top_expert, ne) +# padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) +# padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# +# def benchmark(): +# return ops.padded_gather(x, indices, bin_ids, bins, padded_bins) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('PaddedGather', arguments, mean_t, std_t) +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testPaddedScatter(self, sl, hs, ne): +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# +# # Randomly assign tokens to experts. +# top_expert = torch.randint(0, ne, (sl,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(top_expert, ne) +# padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) +# padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins) +# +# def benchmark(): +# return ops.padded_scatter(x, indices, bin_ids, bins, padded_bins) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('PaddedScatter', arguments, mean_t, std_t) +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testCopy(self, sl, hs, ne): +# # NOTE: Capacity factor == 1. +# # ec = sl // ne +# +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# y = x.clone() +# +# def benchmark(): +# return y.copy_(x) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('Copy', arguments, mean_t, std_t) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/repeat.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/repeat.py new file mode 100644 index 0000000000000000000000000000000000000000..7e9e09de5f857d51cd758ab30b2f3a846d6f9275 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/repeat.py @@ -0,0 +1,10 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch + + +def repeat(x: torch.Tensor, tiling: torch.Size): + if all((t == 1 for t in tiling)): + return x + return x.repeat(*tiling) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/replicate.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/replicate.py new file mode 100644 index 0000000000000000000000000000000000000000..26daf0eede330603a4b8ea7167faf1411d07ca93 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/replicate.py @@ -0,0 +1,36 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + + +# Autograd wrapper for replicate kernel. +class ReplicateOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, bins: torch.Tensor, num_outputs: int): + ctx.save_for_backward(bins) + out = torch.empty((x.shape[0], num_outputs), dtype=x.dtype, device=x.device) + ops.replicate_forward(x, bins, out) + return out + + @staticmethod + def backward(ctx: Any, grad: torch.Tensor): + bins, = ctx.saved_tensors + out = torch.empty((grad.shape[0], bins.shape[0]), dtype=grad.dtype, device=grad.device) + ops.replicate_backward(grad, bins, out) + return out, None, None + + +replicate = ReplicateOp.apply diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/round_up.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/round_up.py new file mode 100644 index 0000000000000000000000000000000000000000..6cf6bc873c9f448c5fa9126ebcfd66e8688002af --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/round_up.py @@ -0,0 +1,14 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch + + +def round_up(x: torch.Tensor, value: int): + assert isinstance(value, int) + assert x.dtype == torch.int32 + + # TODO(tgale): If this becomes and issue + # do this in a custom kernel. We only expect + # to use this on arrays of less than 1k elements. + return torch.div(x + (value - 1), value, rounding_mode='trunc') * value diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/scatter.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/scatter.py new file mode 100644 index 0000000000000000000000000000000000000000..f4605d9b46f387761b070352365f223dbfe69d47 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/scatter.py @@ -0,0 +1,72 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Optional + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for scatter kernel. +class ScatterOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + top_k: int, + ) -> torch.Tensor: + maybe_x = [x] if ctx.needs_input_grad[3] else [] + ctx.save_for_backward(indices, bin_ids, weights, bins, *maybe_x) + ctx.top_k = top_k + ctx.x_shape = x.shape + return kernels.scatter(x, indices, bin_ids, weights, bins, top_k) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + saved_tensors = ctx.saved_tensors + + indices, bin_ids, weights, bins = saved_tensors[:4] + dgrad = None + if ctx.needs_input_grad[0]: + dgrad = kernels.gather( + grad, + indices, + bin_ids, + weights, + bins, + ctx.top_k, + ) + + wgrad = None + if ctx.needs_input_grad[3]: # need wgrad + x = saved_tensors[-1] + wgrad = kernels.scatter_wgrad( + x, + grad, + indices, + bin_ids, + bins, + ctx.top_k, + ) + return dgrad, None, None, wgrad, None, None, None + + +def scatter( + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + top_k: int, +) -> Optional[torch.Tensor]: + return ScatterOp.apply(x, indices, bin_ids, weights, bins, top_k) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/sort.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/sort.py new file mode 100644 index 0000000000000000000000000000000000000000..bda3bf64283e39533c2eae3627e76bb2d0262c9f --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/sort.py @@ -0,0 +1,38 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Optional, Tuple + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + +_BITS_FOR_DTYPE = { + torch.int16: 16, + torch.int32: 32, + torch.int64: 64, +} + + +# Autograd wrapper for sort kernel. +# NOTE: Does not support gradients. +class SortOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, end_bit: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]: + if end_bit is None: + end_bit = _BITS_FOR_DTYPE[x.dtype] + x_out = torch.empty_like(x) + iota_out = torch.empty_like(x) + ops.sort(x, end_bit, x_out, iota_out) + return (x_out, iota_out) + + +sort = SortOp.apply diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/sort_benchmark.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/sort_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..11043c0824c36372585f1d9f48480c2a6ef32eb6 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/sort_benchmark.py @@ -0,0 +1,85 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import numpy as np +import torch +# from absl.testing import parameterized + +from .. import ops + +_SORT_TESTS = ( + (16384, torch.int32, None), + (16384, torch.int32, 2), + (16384, torch.int32, 128), +) + +_BASELINE_SORT_TESTS = ((16384,),) + + +def numpy_dtype(dtype): + types = { + torch.int16: np.int16, + torch.int32: np.int32, + torch.int64: np.int64, + } + return types[dtype] + + +def benchmark_function(fn, iterations=10): + # Run once to get rid of startup overhead. + fn() + times = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + times = np.array(times) + return times.mean(), times.std(), times.max(), times.min() + + +def log_benchmark(arguments, mean_t, std_t): + print('=' * 60) + print('Benchmark Parameters:') + for (key, value) in arguments.items(): + print(f'{key} = {value}') + print('Results:') + print('mean / std = {:.2f}ms / {:.2f}ms'.format(mean_t, std_t)) + print('=' * 60) + + +# class SortBenchmark(parameterized.TestCase): +# +# @parameterized.parameters(*_SORT_TESTS) +# def testSort(self, n, dtype, max_val): +# if max_val is None: +# max_val = np.iinfo(numpy_dtype(dtype)).max +# end_bit = int(np.ceil(np.log2(max_val))) +# x = torch.randint(0, max_val, (n,)).cuda().to(dtype) +# +# mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.sort(x, end_bit),) +# arguments = { +# 'n': n, +# 'dtype': dtype, +# 'max_val': max_val, +# } +# log_benchmark(arguments, mean_t, std_t) +# +# @parameterized.parameters(*_BASELINE_SORT_TESTS) +# def testTorchSort(self, n): +# x = torch.randint(0, 128, (n,)).cuda().to(torch.int32) +# +# mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.sort(x)) +# arguments = { +# 'n': n, +# } +# log_benchmark(arguments, mean_t, std_t) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/stk_autocast.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/stk_autocast.py new file mode 100644 index 0000000000000000000000000000000000000000..7a3626e5e0eec51339c95a448bca84be14a2ca93 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/stk_autocast.py @@ -0,0 +1,39 @@ +# vendored from +# https://github.com/stanford-futuredata/stk/blob/736313768ef697ce13a0594a41b2512a0fbc9884/stk/backend/autocast.py +import functools +import torch + + +def _is_eligible(x): + return x.is_floating_point() and x.is_cuda and (x.dtype is not torch.float64) + + +def _cast(x, dtype): + if isinstance(x, torch.Tensor) and _is_eligible(x): + return x.to(dtype) + elif isinstance(x, map): + return {_cast(k, dtype): _cast(v, dtype) for k, v in x.items()} + elif isinstance(x, list) or isinstance(x, tuple): + return type(x)(map(lambda y: _cast(y, dtype), x)) + return x + + +def custom_fwd(fwd): + """Wrap a custom autograd function that always uses autocast dtype.""" + + @functools.wraps(fwd) + def decorate_fwd(*args, **kwargs): + if torch.is_autocast_enabled(): + with torch.autocast(device_type="cuda", enabled=False): + dtype = torch.get_autocast_gpu_dtype() + return fwd(*_cast(args, dtype), **_cast(kwargs, dtype)) + return fwd(*args, **kwargs) + return decorate_fwd + + +def custom_bwd(bwd): + @functools.wraps(bwd) + def decorate_bwd(*args, **kwargs): + with torch.autocast(device_type="cuda", enabled=False): + return bwd(*args, **kwargs) + return decorate_bwd \ No newline at end of file diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/sum.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/sum.py new file mode 100644 index 0000000000000000000000000000000000000000..e00c1aa68e1193f5b72f75a2edc37de8d505facc --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/sum.py @@ -0,0 +1,9 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +import torch + + +def sum(x: torch.Tensor, dim: int = 0): + if x.shape[dim] == 1: + return x.squeeze(dim=dim) + return x.sum(dim=dim) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/ops/topology.py b/build/torch214-cxx11-cu130-x86_64-linux/ops/topology.py new file mode 100644 index 0000000000000000000000000000000000000000..76a50d3164db20534b099dcb4d8487a7aef25d15 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/ops/topology.py @@ -0,0 +1,45 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + + +# Autograd wrapper for topology kernel. +# NOTE: Does not support gradients. +class TopologyOp(torch.autograd.Function): + + @staticmethod + def forward( + ctx: Any, + padded_bins: torch.Tensor, + block_size: int, + output_block_rows: int, + output_block_columns: int, + ): + out = torch.empty( + output_block_rows * output_block_columns, + dtype=torch.int16, + device=padded_bins.device, + ) + ops.indices( + padded_bins, + block_size, + output_block_rows, + output_block_columns, + out, + ) + return out + + +topology = TopologyOp.apply diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..73c40c267b1c3f4949e9c957a5d2c9f682dfc1a6 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/__init__.py @@ -0,0 +1,7 @@ +# import stk.random +# import stk.ops +# from stk.matrix import Matrix + +from . import random +from . import ops +from .matrix import Matrix diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/backend/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/backend/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/backend/autocast.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/backend/autocast.py new file mode 100644 index 0000000000000000000000000000000000000000..97f6e919a60f3fd579ed0215031008d14111dc96 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/backend/autocast.py @@ -0,0 +1,37 @@ +import functools +import torch + + +def _is_eligible(x): + return x.is_floating_point() and x.is_cuda and (x.dtype is not torch.float64) + + +def _cast(x, dtype): + if isinstance(x, torch.Tensor) and _is_eligible(x): + return x.to(dtype) + elif isinstance(x, map): + return {_cast(k, dtype): _cast(v, dtype) for k, v in x.items()} + elif isinstance(x, list) or isinstance(x, tuple): + return type(x)(map(lambda y: _cast(y, dtype), x)) + return x + + +def custom_fwd(fwd): + """Wrap a custom autograd function that always uses autocast dtype.""" + + @functools.wraps(fwd) + def decorate_fwd(*args, **kwargs): + if torch.is_autocast_enabled(): + with torch.autocast(device_type="cuda", enabled=False): + dtype = torch.get_autocast_gpu_dtype() + return fwd(*_cast(args, dtype), **_cast(kwargs, dtype)) + return fwd(*args, **kwargs) + return decorate_fwd + + +def custom_bwd(bwd): + @functools.wraps(bwd) + def decorate_bwd(*args, **kwargs): + with torch.autocast(device_type="cuda", enabled=False): + return bwd(*args, **kwargs) + return decorate_bwd diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/backend/sputnik.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/backend/sputnik.py new file mode 100644 index 0000000000000000000000000000000000000000..220c947bc1e932e8c77cc30f4069e9930f1aa962 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/backend/sputnik.py @@ -0,0 +1,316 @@ +import torch + +from ..backend import triton_kernels as backend +from ..backend.autocast import custom_bwd, custom_fwd + + +def _standardize_shape(x, transpose): + if transpose: + return torch.Size((x[1], x[0])) + return x + + +def _sparse_transpose(x): + return (torch.Size((x[0][1], x[0][0])), ) + x[1:] + + +def _transpose_helper(x, transpose): + if isinstance(x, torch.Tensor): + return x.t() if transpose else x + if transpose: + x = _sparse_transpose(x) + return x + (transpose,) + + +def _wrap(x): + if isinstance(x, torch.Tensor): + return (x,) + return x + + +def _is_transposed(x): + return (not x.is_contiguous() and + x.stride()[0] == 1 and + x.stride()[1] == x.size()[0]) + + +def _call_helper(op, out, a, b, trans_a, trans_b): + args = (_wrap(_transpose_helper(a, trans_a)) + + _wrap(_transpose_helper(b, trans_b))) + if isinstance(out, tuple): + args = args + out + return op(*args) + + +def _preprocess_inputs(lhs, rhs, dy): + if isinstance(lhs, torch.Tensor) and _is_transposed(lhs): + lhs = lhs.t() + if isinstance(rhs, torch.Tensor) and _is_transposed(rhs): + rhs = rhs.t() + if (isinstance(dy, torch.Tensor) and + not dy.is_contiguous() and + not _is_transposed(dy)): + dy = dy.contiguous() + if isinstance(dy, tuple) and not dy[1].is_contiguous(): + dy = (dy[0], dy[1].contiguous()) + dy[2:] + return lhs, rhs, dy + + +def _postprocess_outputs(x, transpose, grad): + if isinstance(x, torch.Tensor) and transpose: + return grad.t() + return grad + + +def _lhs_gradient(op, lhs, rhs, dy, trans_lhs, trans_rhs): + lhs, rhs, dy = _preprocess_inputs(lhs, rhs, dy) + + a, b = (rhs, dy) if trans_lhs else (dy, rhs) + trans_a = trans_lhs and trans_rhs + trans_b = trans_lhs or not trans_rhs + out = _call_helper(op, lhs, a, b, trans_a, trans_b) + return _postprocess_outputs(lhs, trans_lhs, out) + + +def _rhs_gradient(op, lhs, rhs, dy, trans_lhs, trans_rhs): + lhs, rhs, dy = _preprocess_inputs(lhs, rhs, dy) + + a, b = (dy, lhs) if trans_rhs else (lhs, dy) + trans_a = not trans_lhs or trans_rhs + trans_b = trans_lhs and trans_rhs + out = _call_helper(op, rhs, a, b, trans_a, trans_b) + return _postprocess_outputs(rhs, trans_rhs, out) + + +class DSD(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward(ctx, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_a, + rhs): + ctx.save_for_backward(data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + rhs) + ctx.shape = _standardize_shape(shape, transpose_a) + ctx.transpose_a = transpose_a + + out = torch.empty( + (shape[0], rhs.size()[1]), + dtype=rhs.dtype, + device=rhs.device) + + backend.dsd(shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_a, + rhs, + out) + return out + + @staticmethod + @custom_bwd + def backward(ctx, dy): + saved_tensors = ctx.saved_tensors + lhs = (ctx.shape,) + saved_tensors[:-1] + rhs = saved_tensors[-1] + trans_a = ctx.transpose_a + trans_b = _is_transposed(rhs) + + ddata = None + if ctx.needs_input_grad[1]: + ddata = _lhs_gradient(sdd, + lhs, + rhs, + dy, + trans_a, + trans_b) + drhs = None + if ctx.needs_input_grad[-1]: + op = dds if trans_b else dsd + drhs = _rhs_gradient(op, + lhs, + rhs, + dy, + trans_a, + trans_b) + return None, ddata, None, None, None, None, None, None, None, drhs + + +dsd = DSD.apply + + +class DDS(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward(ctx, + lhs, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_b): + ctx.save_for_backward(lhs, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t) + ctx.shape = _standardize_shape(shape, transpose_b) + ctx.transpose_b = transpose_b + out = torch.empty((lhs.size()[0], shape[1]), + dtype=lhs.dtype, + device=lhs.device) + backend.dds(lhs, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_b, + out) + return out + + @staticmethod + @custom_bwd + def backward(ctx, dy): + saved_tensors = ctx.saved_tensors + lhs = saved_tensors[0] + rhs = (ctx.shape,) + saved_tensors[1:] + trans_a = _is_transposed(lhs) + trans_b = ctx.transpose_b + + dlhs = None + if ctx.needs_input_grad[0]: + op = dsd if trans_a else dds + dlhs = _lhs_gradient(op, + lhs, + rhs, + dy, + trans_a, + trans_b) + ddata = None + if ctx.needs_input_grad[2]: + ddata = _rhs_gradient(sdd, + lhs, + rhs, + dy, + trans_a, + trans_b) + return dlhs, None, ddata, None, None, None, None, None, None, None + + +dds = DDS.apply + + +class SDD(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward(ctx, + lhs, + rhs, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t): + ctx.save_for_backward( + lhs, + rhs, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t) + ctx.shape = shape + out = torch.empty( + data.shape, + dtype=lhs.dtype, + device=lhs.device) + backend.sdd(lhs, + rhs, + shape, + out, + offsets, + row_indices, + column_indices) + return out + + @staticmethod + @custom_bwd + def backward(ctx, dy): + saved_tensors = ctx.saved_tensors + lhs, rhs = saved_tensors[:2] + dy = (ctx.shape, dy) + saved_tensors[2:] + trans_a = _is_transposed(lhs) + trans_b = _is_transposed(rhs) + + dlhs = None + if ctx.needs_input_grad[0]: + op = dds if trans_a else dsd + dlhs = _lhs_gradient(op, + lhs, + rhs, + dy, + trans_a, + trans_b) + drhs = None + if ctx.needs_input_grad[1]: + op = dsd if trans_b else dds + drhs = _rhs_gradient(op, + lhs, + rhs, + dy, + trans_a, + trans_b) + return dlhs, drhs, None, None, None, None, None, None, None, None + + +sdd = SDD.apply + +class RowIndices(torch.autograd.Function): + + @staticmethod + def forward(ctx, shape, data, offsets, column_indices): + out = torch.empty( + column_indices.shape, + dtype=column_indices.dtype, + device=column_indices.device) + backend.row_indices(shape, data, offsets, column_indices, out) + return out + + +row_indices = RowIndices.apply diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/backend/triton_kernels.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/backend/triton_kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..c535309f3321249f475367164a558f94a4f8eb86 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/backend/triton_kernels.py @@ -0,0 +1,393 @@ +import torch +import triton +import triton.language as tl +from dataclasses import dataclass + +@dataclass +class TritonConfig: + BLOCK_M: int = 128 + BLOCK_N: int = 128 + BLOCK_K: int = 32 + BLOCK_SIZE: int = 128 + NUM_STAGES: int = 4 + NUM_WARPS: int = 4 + +def _validate_matmul_dims(M: int, K: int, N: int): + error_string = "incompatible dimensions: tensor has dim with length: {}, which must be divisible by {}" + assert M % TritonConfig.BLOCK_M == 0, error_string.format(M, TritonConfig.BLOCK_M) + assert K % TritonConfig.BLOCK_K == 0, error_string.format(K, TritonConfig.BLOCK_K) + assert N % TritonConfig.BLOCK_N == 0, error_string.format(N, TritonConfig.BLOCK_N) + +@triton.autotune( + configs=[ + # basic configs for compute-bound matmuls + triton.Config({ + 'BLOCK_M': TritonConfig.BLOCK_M, + 'BLOCK_N': TritonConfig.BLOCK_N, + 'BLOCK_K': TritonConfig.BLOCK_K, + 'BLOCK_SIZE': TritonConfig.BLOCK_SIZE + }, num_stages=TritonConfig.NUM_STAGES, num_warps=TritonConfig.NUM_WARPS), + ], + key=['M', 'N', 'K'], +) +@triton.jit +def _sdd_kernel(A, B, C, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + row_indices, column_indices, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + BLOCK_SIZE: tl.constexpr, GROUP_M: tl.constexpr, ACC_TYPE: tl.constexpr, + ): + # matrix multiplication + pid = tl.program_id(0) + pid_m = tl.load(row_indices + pid) + pid_n = tl.load(column_indices + pid) + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) + rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) + rk = tl.arange(0, BLOCK_K) + # pointers + A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) + B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) + # do matrix multiplication + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) + for k in range(0, tl.cdiv(K, BLOCK_K)): + a = tl.load(A) + b = tl.load(B) + acc += tl.dot(a, b) + A += BLOCK_K * stride_ak + B += BLOCK_K * stride_bk + #Store to sparse matrix + acc = acc.to(C.dtype.element_ty) + BLOCK_ELEMENTS = BLOCK_SIZE * BLOCK_SIZE + cm = tl.arange(0, BLOCK_M) + cn = tl.arange(0, BLOCK_N) + C = C + pid * BLOCK_ELEMENTS + (cm[:, None] * stride_cm + cn[None, :] * stride_cn) + tl.store(C, acc, mask=True) + +@triton.autotune( + configs=[ + # basic configs for compute-bound matmuls + triton.Config({ + 'BLOCK_M': TritonConfig.BLOCK_M, + 'BLOCK_N': TritonConfig.BLOCK_N, + 'BLOCK_K': TritonConfig.BLOCK_K, + 'BLOCK_SIZE': TritonConfig.BLOCK_SIZE + }, num_stages=TritonConfig.NUM_STAGES, num_warps=TritonConfig.NUM_WARPS), + ], + key=['M', 'N', 'K'], +) +@triton.jit +def _dsd_kernel(A, B, C, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + row_indices, column_indices, offsets, + block_offsets_t, trans_A: tl.constexpr, trans_B: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + BLOCK_SIZE: tl.constexpr, GROUP_M: tl.constexpr, ACC_TYPE: tl.constexpr, + ): + + # matrix multiplication + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + num_pid_m = tl.num_programs(0) + num_pid_n = tl.num_programs(1) + pid_n, pid_m = tl.swizzle2d(pid_n, pid_m, num_pid_n, num_pid_m, GROUP_M) + + start_inx = tl.load(offsets + pid_m) + end_inx = tl.load(offsets + pid_m + 1) + + # pointers to sparse matrix + rm = tl.arange(0, BLOCK_M) + rak = tl.arange(0, BLOCK_K) + + A += (rm[:, None] * stride_am + rak[None, :] * stride_ak) + + # pointers to dense matrix + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + rbk = tl.arange(0, BLOCK_K) + B += (rbk[:, None] * stride_bk + rn[None, :] * stride_bn) + + # do matrix multiplication + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) + nsub_blocks = tl.cdiv(BLOCK_SIZE, BLOCK_K) + + BLOCK_ELEMENTS = BLOCK_SIZE * BLOCK_SIZE + ak_sub_incr = BLOCK_K * stride_ak + bk_sub_incr = BLOCK_K * stride_bk + bk_block_incr = BLOCK_SIZE * stride_bk + + for k in range(nsub_blocks * (end_inx - start_inx)): + sub_block_inx = k % nsub_blocks + block_inx = k // nsub_blocks + + if trans_A: + ptr_A = A + tl.load(block_offsets_t + start_inx + block_inx) * BLOCK_ELEMENTS + sub_block_inx * ak_sub_incr + else: + ptr_A = A + (start_inx + block_inx) * BLOCK_ELEMENTS + sub_block_inx * ak_sub_incr + + ptr_B = B + tl.load(column_indices + start_inx + block_inx) * bk_block_incr + sub_block_inx * bk_sub_incr + + a = tl.load(ptr_A) + b = tl.load(ptr_B) + acc += tl.dot(a, b) + + acc = acc.to(C.dtype.element_ty) + + cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + C = C + (cm[:, None] * stride_cm + cn[None, :] * stride_cn) + tl.store(C, acc, mask=True) + +@triton.autotune( + configs=[ + # basic configs for compute-bound matmuls + triton.Config({ + 'BLOCK_M': TritonConfig.BLOCK_M, + 'BLOCK_N': TritonConfig.BLOCK_N, + 'BLOCK_K': TritonConfig.BLOCK_K, + 'BLOCK_SIZE': TritonConfig.BLOCK_SIZE + }, num_stages=TritonConfig.NUM_STAGES, num_warps=TritonConfig.NUM_WARPS), + ], + key=['M', 'N', 'K'], +) +@triton.jit +def _dds_kernel(A, B, C, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + row_indices, column_indices, offsets, + block_offsets_t, trans_A: tl.constexpr, trans_B: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + BLOCK_SIZE: tl.constexpr, GROUP_M: tl.constexpr, ACC_TYPE: tl.constexpr, + ): + + # matrix multiplication + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + num_pid_m = tl.num_programs(0) + num_pid_n = tl.num_programs(1) + pid_n, pid_m = tl.swizzle2d(pid_n, pid_m, num_pid_n, num_pid_m, GROUP_M) + + start_inx = tl.load(offsets + pid_n) + end_inx = tl.load(offsets + pid_n + 1) + + # pointers to dense matrix + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rak = tl.arange(0, BLOCK_K) + + A += (rm[:, None] * stride_am + rak[None, :] * stride_ak) + + # pointers to sparse matrix + rn = tl.arange(0, BLOCK_N) + rbk = tl.arange(0, BLOCK_K) + B += (rbk[:, None] * stride_bk + rn[None, :] * stride_bn) + + # do matrix multiplication + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) + nsub_blocks = tl.cdiv(BLOCK_SIZE, BLOCK_K) + + BLOCK_ELEMENTS = BLOCK_SIZE * BLOCK_SIZE + + ak_sub_incr = BLOCK_K * stride_ak + ak_block_incr = BLOCK_SIZE * stride_ak + bk_sub_incr = BLOCK_K * stride_bk + + for k in range(nsub_blocks * (end_inx - start_inx)): + sub_block_inx = k % nsub_blocks + block_inx = k // nsub_blocks + + if trans_B: + ptr_B = B + (start_inx + block_inx) * BLOCK_ELEMENTS + sub_block_inx * bk_sub_incr + else: + ptr_B = B + tl.load(block_offsets_t + start_inx + block_inx) * BLOCK_ELEMENTS + sub_block_inx * bk_sub_incr + + ptr_A = A + tl.load(column_indices + start_inx + block_inx) * ak_block_incr + sub_block_inx * ak_sub_incr + a = tl.load(ptr_A) + b = tl.load(ptr_B) + acc += tl.dot(a, b) + + acc = acc.to(C.dtype.element_ty) + cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + C = C + (cm[:, None] * stride_cm + cn[None, :] * stride_cn) + tl.store(C, acc, mask=True) + +def dsd(shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_a, + rhs, + out + ): + + device = rhs.device + trans_A = transpose_a + trans_B = False + + if rhs.stride(0) > 1 and rhs.stride(1) > 1: + trans_B = True + + # checks constraints + assert shape[1] == rhs.shape[0], "incompatible dimensions" + M, K = shape + _, N = rhs.shape + + _validate_matmul_dims(M, K, N) + + # accumulator types + ACC_TYPE = tl.float32 if rhs.dtype in [torch.float16, torch.bfloat16, torch.float32] else tl.int32 + + stride_am, stride_ak = data.stride(1), data.stride(2) + stride_bk, stride_bn = rhs.stride(0), rhs.stride(1) + a_column_indices = column_indices + a_offsets = offsets + + # launch kernel + grid = lambda META: (triton.cdiv(M, META['BLOCK_M']), triton.cdiv(N, META['BLOCK_N'])) + + if trans_A: + stride_am, stride_ak = data.stride(2), data.stride(1) + a_column_indices, a_offsets = column_indices_t, offsets_t + + if trans_B: + stride_bk, stride_bn = rhs.stride(1), rhs.stride(0) + + _dsd_kernel[grid]( + data.data, rhs, out, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + out.stride(0), out.stride(1), + row_indices, a_column_indices, a_offsets, + block_offsets_t, trans_A, trans_B, + GROUP_M=128, ACC_TYPE=ACC_TYPE + ) + # return out + +def dds(lhs, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_b, + out + ): + + device = lhs.device + trans_B = transpose_b + trans_A = False + + if lhs.stride(0) > 1 and lhs.stride(1) > 1: + trans_A = True + + # checks constraints + assert lhs.shape[1] == shape[0], "incompatible dimensions" + M, K = lhs.shape + _, N = shape + + _validate_matmul_dims(M, K, N) + + # accumulator types + ACC_TYPE = tl.float32 if lhs.dtype in [torch.float16, torch.bfloat16, torch.float32] else tl.int32 + + stride_am, stride_ak = lhs.stride(0), lhs.stride(1) + stride_bk, stride_bn = data.stride(1), data.stride(2) + b_column_indices = column_indices_t + b_offsets = offsets_t + + # launch kernel + grid = lambda META: (triton.cdiv(M, META['BLOCK_M']), triton.cdiv(N, META['BLOCK_N'])) + + if trans_A: + stride_am, stride_ak = lhs.stride(1), lhs.stride(0) + if trans_B: + stride_bk, stride_bn = data.stride(2), data.stride(1) + b_column_indices, b_offsets = column_indices, offsets + + _dds_kernel[grid]( + lhs, data, out, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + out.stride(0), out.stride(1), + row_indices, b_column_indices, b_offsets, + block_offsets_t, trans_A, trans_B, + GROUP_M=128, ACC_TYPE=ACC_TYPE + ) + +def sdd(lhs, + rhs, + shape, + out, + offsets, + row_indices, + column_indices + ): + + device = out.device + trans_A = False + trans_B = False + + if lhs.stride(0) > 1 and lhs.stride(1) > 1: + trans_A = True + if rhs.stride(0) > 1 and rhs.stride(1) > 1: + trans_B = True + + # checks constraints + assert lhs.shape[1] == rhs.shape[0], "incompatible dimensions" + M, K = lhs.shape + _, N = rhs.shape + + _validate_matmul_dims(M, K, N) + + # accumulator types + ACC_TYPE = tl.float32 if out.dtype in [torch.float16, torch.bfloat16, torch.float32] else tl.int32 + + # launch kernel + nnz_blocks = len(row_indices) + grid = lambda META: (nnz_blocks,) + + stride_am, stride_ak = lhs.stride(0), lhs.stride(1) + stride_bk, stride_bn = rhs.stride(0), rhs.stride(1) + + if trans_A: + stride_am, stride_ak = lhs.stride(1), lhs.stride(0) + if trans_B: + stride_bk, stride_bn = rhs.stride(1), rhs.stride(0) + + _sdd_kernel[grid]( + lhs, rhs, out, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + out.stride(1), out.stride(2), + row_indices, column_indices, + GROUP_M=128, ACC_TYPE=ACC_TYPE + ) + +@triton.jit +def _row_indices_kernel(offsets, out): + pid = tl.program_id(0) + row_offset = tl.load(offsets + pid) + nnz_blocks = tl.load(offsets + pid + 1) - row_offset + for nnz_block in range(nnz_blocks): + tl.store(out + row_offset + nnz_block, pid) + +def row_indices( + shape, data, offsets, column_indices, out +): + block_rows = len(offsets) - 1 + _row_indices_kernel[(block_rows, )](offsets, out) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/matrix.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/matrix.py new file mode 100644 index 0000000000000000000000000000000000000000..80f42263d6aada287adbfa52a61fe950162a9e28 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/matrix.py @@ -0,0 +1,329 @@ +import numpy as np +import torch + +# 1. Add heavyweight (data) validation helper. +# 2. Add construction helpers +# 3. Make indentation consistent +# 4. Replace asserts with descriptive errors. + +## +### Validation helpers. +## + + +def _validate_matrix(shape, data, row_indices, column_indices, offsets): + # Data should be [nnz, block_size, block_size] + if data.dim() == 1: + data = torch.reshape(data, [data.numel(), 1, 1]) + + # Blocks should be square. + if data.shape[-2] != data.shape[-1]: + raise ValueError( + "Expected square blocking in data. " + f"Got block shape {[data.shape[-2], data.shape[-1]]}") + + # Flatten batch dimensions on data - original shape preserved + # in shape argument. + block_size = data.shape[-1] + data = data.view([-1, block_size, block_size]) + + if data.dim() != 3: + raise ValueError( + "Expected 3D shape for data (nnz, block, block). " + f"Got shape {data.dim()}D shape.") + + block_size = data.shape[1] + if shape[-2] % block_size != 0 or shape[-1] % block_size != 0: + raise ValueError( + "Matrix shape must be dividible by blocking. " + f"Got shape {shape} with " + f"{[block_size, block_size]} blocking.") + + if np.prod(shape) < data.numel(): + raise ValueError( + "Invalid matrix. Number of nonzeros exceeds matrix capacity " + f"({data.numel()} v. {np.prod(shape)})") + + if row_indices.dim() != 1: + raise ValueError( + f"Expected 1D row_indices. Got {row_indices.dim()}D row_indices.") + + if column_indices.dim() != 1: + raise ValueError( + f"Expected 1D column_indices. Got {column_indices.dim()}D column_indices.") + + if offsets.dim() != 1: + raise ValueError( + f"Expected 1D offsets. Got {offsets.dim()}D offsets.") + + if row_indices.numel() != data.shape[0]: + raise ValueError( + "Expected 1 index per nonzero block. " + f"Got {row_indices.numel()} row_indices for {data.shape[0]} blocks") + + if column_indices.numel() != data.shape[0]: + raise ValueError( + "Expected 1 index per nonzero block. " + f"Got {column_indices.numel()} column_indices for {data.shape[0]} blocks") + + block_rows = np.prod(shape[:-1]) / block_size + if offsets.numel() != block_rows + 1: + raise ValueError( + "Expected one offset per block row plus one. " + f"Got {offsets.numel()} offsets with {block_rows} block rows.") + + is_cuda = (data.is_cuda and + row_indices.is_cuda and + column_indices.is_cuda and + offsets.is_cuda) + is_cpu = (not data.is_cuda and + not row_indices.is_cuda and + not column_indices.is_cuda and + not offsets.is_cuda) + if not (is_cuda or is_cpu): + raise ValueError( + "Expected data & meta-data on common device. " + f"Got data on {data.device}, row_indices on {row_indices.device} " + f"column_indices on {column_indices.device} and " + f"offsets on {offsets.device}.") + + if data.dtype != torch.float16: + raise ValueError( + f"Expected float16 data. Got {data.dtype} data.") + if row_indices.dtype != torch.int16: + raise ValueError( + f"Expected int16 row_indices. Got {row_indices.dtype} row_indices.") + if column_indices.dtype != torch.int16: + raise ValueError( + f"Expected int16 column_indices. Got {column_indices.dtype} column_indices.") + if offsets.dtype != torch.int32: + raise ValueError( + f"Expected int32 offsets. Got {offsets.dtype} offsets.") + return data + + +def _transpose(size, data, row_indices, column_indices, offsets): + block_columns = size[1] // data.shape[1] + + # Sort row indices by column indices to get the transposed matrix's + # column indices. + gather_indices = column_indices.argsort() + column_indices_t = row_indices.gather(0, gather_indices) + block_offsets_t = gather_indices.int() + + # NOTE: Histogram is not implemented for any integer type on CPU. Do + # the histogram in 32-bit float, which can exactly represent 16-bit + # integers. + column_indices_float = column_indices.float() + + zero = torch.zeros((1,), dtype=torch.int32, device=data.device) + nnz_per_column = column_indices_float.histc(block_columns, 0, block_columns) + nnz_per_column = nnz_per_column.int() + offsets_t = torch.cat([zero, nnz_per_column.cumsum(0, dtype=torch.int32)]) + return column_indices_t, offsets_t, block_offsets_t + + +class Matrix(torch.nn.Module): + """A matrix stored in sparse format. + + Underlying format is block compressed sparse row (BCSR). + + TODO(tgale): Make this mirror torch.Tensor API as much as possible. + """ + + def __init__(self, + size, + data, + row_indices, + column_indices, + offsets, + column_indices_t=None, + offsets_t=None, + block_offsets_t=None): + super().__init__() + self._size = size + self._data = data + self._row_indices = row_indices + self._column_indices = column_indices + self._offsets = offsets + + # Produce the transpose meta-data if it is not passed in. + if ((column_indices_t is None) or (offsets_t is None) or + (block_offsets_t is None)): + column_indices_t, offsets_t, block_offsets_t = _transpose( + size, data, row_indices, column_indices, offsets) + self._column_indices_t = column_indices_t + self._offsets_t = offsets_t + self._block_offsets_t = block_offsets_t + + self._transposed = False + + # Validate that our metadata will not overflow. + max_dim = np.iinfo(np.int16).max * self.blocking + if column_indices.dtype == torch.int16: + if size[0] > max_dim or size[1] > max_dim: + raise ValueError( + "Sparse matrix with shape {size} exceeds representable " + "size with 16-bit indices.") + + def validate(self): + _validate_matrix(self._size, + self._data, + self._row_indices, + self._column_indices, + self._offsets) + + # TODO(tgale): Add heavyweight data validation. + + def to(self, device): + # TODO(tgale): Handle type conversions here. We + # need to set the appropriate meta-data type for + # the given floating-point type. + self._data = self._data.to(device) + self._row_indices = self._row_indices.to(device) + self._column_indices = self._column_indices.to(device) + self._offsets = self._offsets.to(device) + self._column_indices_t = self._column_indices_t.to(device) + self._offsets_t = self._offsets_t.to(device) + self._block_offsets_t = self._block_offsets_t.to(device) + return self + + def cuda(self): + return self.to(torch.cuda.current_device()) + + def clone(self): + return Matrix( + self.size(), + self.data.clone(), + self.row_indices.clone(), + self.column_indices.clone(), + self.offsets.clone(), + self.column_indices_t.clone(), + self.offsets_t.clone(), + self.block_offsets_t.clone()) + + def t(self): + if self.dim() != 2: + raise ValueError( + "t() expects a tensor with <= 2 dimensions, " + f"but self is {self.dim()}D.") + out = Matrix(self.size(), + self.data, + self.row_indices, + self.column_indices, + self.offsets, + self.column_indices_t, + self.offsets_t, + self.block_offsets_t) + out._transposed = not self._transposed + out._size = torch.Size((self._size[1], self._size[0])) + return out + + def contiguous(self): + raise ValueError("Not yet implemented.") + + def is_contiguous(self): + return not self._transposed + + @property + def is_cuda(self): + return self._data.is_cuda + + @property + def device(self): + return self._data.device + + def size(self): + return self._size + + @property + def shape(self): + return self.size() + + def dim(self): + return len(self._size) + + @property + def data(self): + return self._data + + @property + def row_indices(self): + return self._row_indices + + @property + def column_indices(self): + return self._column_indices + + @property + def offsets(self): + return self._offsets + + @property + def offsets_t(self): + return self._offsets_t + + @property + def column_indices_t(self): + return self._column_indices_t + + @property + def block_offsets_t(self): + return self._block_offsets_t + + @property + def dtype(self): + return self.data.dtype + + @property + def nnz(self): + return self.data.numel() + + @property + def blocking(self): + return self.data.shape[1] + + @property + def requires_grad(self): + return self.data.requires_grad + + def requires_grad_(self, x): + self.data.requires_grad_(x) + return self + + def view(self, *shape): + assert self.is_contiguous() + if shape[-1] != self.size()[-1]: + raise ValueError( + "Can't change view on compressed dimension. " + f"{self.size()[-1]} v. {shape[-1]}.") + if np.prod(shape) != np.prod(self.size()): + raise ValueError( + "Mismatch in numel of Matrix and new shape. " + f"{np.prod(self.size())} v. {np.prod(shape)}") + return Matrix(shape, + self.data, + self.row_indices, + self.column_indices, + self.offsets, + self.column_indices_t, + self.offsets_t, + self.block_offsets_t) + + @property + def grad(self): + # TODO(tgale): Make sure this mirrors torch.Tensor + # behavior in the case where we ask for the gradient + # of a non-contiguous tensor. + size = self.size() + if not self.is_contiguous(): + size = torch.Size((size[1], size[0])) + out = Matrix(size, + self.data.grad, + self.row_indices, + self.column_indices, + self.offsets, + self.column_indices_t, + self.offsets_t, + self.block_offsets_t) + return out if self.is_contiguous() else out.t() diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fc873b236f4cd4036964c016a4036e3ce5ebf1ac --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/__init__.py @@ -0,0 +1,3 @@ +from .linear_ops import dds, dsd, sdd +from .matrix_ops import ones_like, row_indices, sum, to_dense, to_sparse +from .eltwise_ops import mul diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/eltwise_ops.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/eltwise_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ba7d7332320250fd01fa60e60528f19de3e8ed03 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/eltwise_ops.py @@ -0,0 +1,28 @@ +from ..matrix import Matrix + +def mul(a, b): + """Performs element-wise multiplication of matrices a and b. + + It is the user's responsibility to make sure that a and b + follow the same matrix topology. This function assumes it is safe + to use the topoplogy of a. + + Args: + a: stk.Matrix. + b: stk.Matrix with a's matrix topology. + + Returns: + stk.Matrix where the entries correspond to torch.mul(a, b). + """ + assert isinstance(a, Matrix) + assert isinstance(b, Matrix) + assert a.size() == b.size() + + return Matrix(a.size(), + a.data * b.data, + a.row_indices, + a.column_indices, + a.offsets, + a.column_indices_t, + a.offsets_t, + a.block_offsets_t) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/eltwise_ops_test.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/eltwise_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2939372a5c68ac92b47b11015db4f75f4fd60ffa --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/eltwise_ops_test.py @@ -0,0 +1,86 @@ +import unittest +import itertools +import torch +# from absl.testing import parameterized + +import stk +from stk.ops.linear_ops_test import allclose, _dense_and_sparse + +_MATRIX_SIZES = ( + (128, 128, 0.0), + (256, 256, 0.5), + (2048, 1024, 0.8), + (512, 128, 0.0), + (128, 512, 0.0), + (1024, 512, 0.0), + (1024, 512, 0.5), + (1024, 512, 0.75), + (512, 1024, 0.0), + (512, 1024, 0.5), + (512, 1024, 0.75), + (1024, 1024, 0.0), + (1024, 1024, 0.5), + (1024, 1024, 0.75), +) + +_DTYPE = ( + torch.float16, torch.bfloat16 +) + +def _generate_testcases(): + testcases = itertools.product(_MATRIX_SIZES, _DTYPE) + testcases = [(*size, 128, dtype) for + (size, dtype) in testcases] + return testcases + +_ELTWISE_OP_TESTS = _generate_testcases() + +def _dense_and_sparse_like(x, std=0.1): + dense_data = torch.randn_like(x.data, device=x.device) * std + sparse = stk.Matrix(x.size(), + dense_data, + x.row_indices, + x.column_indices, + x.offsets) + dense = stk.ops.to_dense(sparse) + + return (dense.requires_grad_(True), + sparse.requires_grad_(True)) + +# @parameterized.parameters(_ELTWISE_OP_TESTS) +# class EltwiseOpsTest(parameterized.TestCase): +# +# def testEltwiseMul(self, m, n, sparsity, blocking, dtype): +# +# a_dense, a = _dense_and_sparse(m, n, sparsity, blocking, dtype) +# b_dense, b = _dense_and_sparse_like(a) +# +# out = stk.ops.mul(a, b) +# expected_out = torch.mul(a_dense, b_dense) +# +# # Compute the gradients w.r.t. the inputs. +# expected_out.sum().backward() +# stk.ops.sum(out).backward() +# +# # Validate the results. +# out = stk.ops.to_dense(out) +# self.assertEqual(out.dim(), 2) +# self.assertEqual(expected_out.size(), out.size()) +# self.assertTrue(allclose(out, expected_out)) +# +# # LHS gradient. +# grad = stk.ops.to_dense(a.grad) +# expected_grad = a_dense.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size(), grad.size()) +# self.assertTrue(allclose(grad, expected_grad)) +# +# # RHS gradient. +# grad = stk.ops.to_dense(b.grad) +# expected_grad = b_dense.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size(), grad.size()) +# self.assertTrue(allclose(grad, expected_grad)) + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/linear_ops.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/linear_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..9d277c8c07f9e30addc31900a12175c8a1f4d7ad --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/linear_ops.py @@ -0,0 +1,59 @@ +import torch + +from ..backend import sputnik +from ..matrix import Matrix + + +def dsd(a, b): + assert isinstance(a, Matrix) + assert isinstance(b, torch.Tensor) + return sputnik.dsd( + a.size(), + a.data, a.offsets, + a.row_indices, + a.column_indices, + a.offsets_t, + a.column_indices_t, + a.block_offsets_t, + not a.is_contiguous(), + b) + + +def dds(a, b): + assert isinstance(a, torch.Tensor) + assert isinstance(b, Matrix) + return sputnik.dds( + a, + b.size(), + b.data, b.offsets, + b.row_indices, + b.column_indices, + b.offsets_t, + b.column_indices_t, + b.block_offsets_t, + not b.is_contiguous()) + + +def sdd(a, b, topo): + assert isinstance(a, torch.Tensor) + assert isinstance(b, torch.Tensor) + assert isinstance(topo, Matrix) + assert topo.is_contiguous() + out = sputnik.sdd( + a, b, + topo.size(), + topo.data, + topo.offsets, + topo.row_indices, + topo.column_indices, + topo.offsets_t, + topo.column_indices_t, + topo.block_offsets_t) + return Matrix(topo.size(), + out, + topo.row_indices, + topo.column_indices, + topo.offsets, + topo.column_indices_t, + topo.offsets_t, + topo.block_offsets_t) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/linear_ops_test.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/linear_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..c1c24d350df9c1b2346c7da885502cd696c88867 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/linear_ops_test.py @@ -0,0 +1,216 @@ +import unittest +import itertools +import numpy as np +import torch +# from absl.testing import parameterized + +import stk + + +def allclose(x, y, pct=0.25): + mask = torch.isclose(x, y, rtol=5e-2) + pct_diff = (mask.numel() - mask.sum()) / mask.numel() * 100 + if pct_diff > pct: + print("{:.2f}% of values not close.".format(pct_diff)) + return False + return True + + +# An assortment of problems designed to make sure +# the bindings are operating correctly. +_MATRIX_SIZES = ( + (128, 128, 128, 0.0), + (256, 256, 256, 0.5), + (2048, 1024, 512, 0.8), + (512, 128, 128, 0.0), + (128, 128, 512, 0.0), + (1024, 512, 512, 0.0), + (1024, 512, 512, 0.5), + (1024, 512, 512, 0.75), + (512, 512, 1024, 0.0), + (512, 512, 1024, 0.5), + (512, 512, 1024, 0.75), + (1024, 1024, 1024, 0.0), + (1024, 1024, 1024, 0.5), + (1024, 1024, 1024, 0.75), +) + +_TRANSPOSE = ( + (False, False), + (False, True), + (True, False), + (True, True), +) + +_DTYPE = ( + torch.float16, torch.bfloat16 +) + +def _generate_testcases(): + testcases = itertools.product(_MATRIX_SIZES, _TRANSPOSE, _DTYPE) + testcases = [(*size, *trans, 128, dtype) for + (size, trans, dtype) in testcases] + return testcases + +_LINEAR_OP_TESTS = _generate_testcases() + +def _dense_and_sparse(rows, cols, sparsity, blocking, dtype, std=0.1): + mask = stk.random.dense_mask(rows, cols, sparsity, blocking) + dense = (torch.randn(rows, cols) * std * mask).type(dtype) + sparse = stk.ops.to_sparse(dense, blocking) + cuda_device = torch.device("cuda") + return (dense.to(cuda_device).requires_grad_(True), + sparse.to(cuda_device).requires_grad_(True)) + + +def _dense(rows, cols, dtype, std=0.1): + cuda_device = torch.device("cuda") + out = (torch.randn(rows, cols) * std).type(dtype) + return out.to(cuda_device).requires_grad_(True) + + +def _dense_2x(rows, cols, dtype): + a = _dense(rows, cols, dtype) + return a, a.detach().requires_grad_(True) + + +def _with_transpose(op, a, b, trans_a, trans_b): + a = a.t() if trans_a else a + b = b.t() if trans_b else b + return op(a, b) + + +def _mmm(a, b, topo): + mask = stk.ops.to_dense(stk.ops.ones_like(topo)) + return torch.mm(a, b) * mask + + +def _sparse_out_with_transpose(op, a, b, topo, trans_a, trans_b): + a = a.t() if trans_a else a + b = b.t() if trans_b else b + return op(a, b, topo) + + +def _mask(x, mask): + mask = stk.ops.to_dense(stk.ops.ones_like(mask)) + return x * mask + + +# @parameterized.parameters(*_LINEAR_OP_TESTS) +# class LinearOpsTest(parameterized.TestCase): +# +# def testLinearOps_Dsd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype): +# # Construct the operands. +# a_shape = (k, m) if trans_a else (m, k) +# a_dense, a = _dense_and_sparse(*a_shape, sparsity, blocking, dtype) +# b_shape = (n, k) if trans_b else (k, n) +# b, bcp = _dense_2x(*b_shape, dtype) +# +# # Execute the matmul. +# out = _with_transpose(stk.ops.dsd, a, b, trans_a, trans_b) +# expected_out = _with_transpose(torch.mm, a_dense, bcp, trans_a, trans_b) +# +# # Compute the gradients w.r.t. the inputs. +# expected_out.sum().backward() +# out.sum().backward() +# +# # Validate the results. +# self.assertEqual(out.dim(), 2) +# self.assertEqual(expected_out.size()[0], out.size()[0]) +# self.assertEqual(expected_out.size()[1], out.size()[1]) +# self.assertTrue(allclose(out, expected_out)) +# +# # LHS gradient. +# grad = stk.ops.to_dense(a.grad) +# expected_grad = _mask(a_dense.grad, a.grad) +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# # RHS gradient. +# grad = b.grad +# expected_grad = bcp.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# def testLinearOps_Dds(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype): +# # Construct the operands. +# a_shape = (k, m) if trans_a else (m, k) +# a, acp = _dense_2x(*a_shape, dtype) +# b_shape = (n, k) if trans_b else (k, n) +# b_dense, b = _dense_and_sparse(*b_shape, sparsity, blocking, dtype) +# +# # Execute the matmul. +# out = _with_transpose(stk.ops.dds, a, b, trans_a, trans_b) +# expected_out = _with_transpose(torch.mm, acp, b_dense, trans_a, trans_b) +# +# # Compute the gradients w.r.t. the inputs. +# expected_out.sum().backward() +# out.sum().backward() +# +# # Validate the results. +# self.assertEqual(out.dim(), 2) +# self.assertEqual(expected_out.size()[0], out.size()[0]) +# self.assertEqual(expected_out.size()[1], out.size()[1]) +# self.assertTrue(allclose(out, expected_out)) +# +# # LHS gradient. +# grad = a.grad +# expected_grad = acp.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# # RHS gradient. +# grad = stk.ops.to_dense(b.grad) +# expected_grad = _mask(b_dense.grad, b.grad) +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# def testLinearOps_Sdd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype): +# # Construct the operands. +# a_shape = (k, m) if trans_a else (m, k) +# a, acp = _dense_2x(*a_shape, dtype) +# b_shape = (n, k) if trans_b else (k, n) +# b, bcp = _dense_2x(*b_shape, dtype) +# _, topo = _dense_and_sparse(m, n, sparsity, blocking, dtype) +# +# # Execute the matmul. +# out = _sparse_out_with_transpose(stk.ops.sdd, a, b, topo, trans_a, trans_b) +# expected_out = _sparse_out_with_transpose(_mmm, acp, bcp, topo, trans_a, trans_b) +# +# # Compute the gradients w.r.t. the inputs. +# expected_out.sum().backward() +# stk.ops.sum(out).backward() +# +# # Validate the results. +# out = stk.ops.to_dense(out) +# self.assertEqual(out.dim(), 2) +# self.assertEqual(expected_out.size()[0], out.size()[0]) +# self.assertEqual(expected_out.size()[1], out.size()[1]) +# self.assertTrue(allclose(out, expected_out)) +# +# # LHS gradient. +# grad = a.grad +# expected_grad = acp.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# # RHS gradient. +# grad = b.grad +# expected_grad = bcp.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/matrix_ops.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/matrix_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..447c72dc73439d84f58c917676cc04e64f13e97d --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/matrix_ops.py @@ -0,0 +1,98 @@ +from ..backend import sputnik +from ..matrix import Matrix +import torch +import numpy as np + + +@torch.no_grad() +def row_indices(shape, data, offsets, column_indices): + return sputnik.row_indices(shape, data, offsets, column_indices) + + +# TODO(tgale): Replace this helper with a custom kernel. This operation +# is much simpler to do than how it's currently implemented. +@torch.no_grad() +def _expand_for_blocking(idxs, blocking): + # Duplicate for block column dimension. + idxs = torch.reshape(idxs, [idxs.size()[0], 1, 2]).repeat(1, blocking, 1) + + # Update the column indices. + idxs[:, :, 1] *= blocking + idxs[:, :, 1] += torch.reshape(torch.arange(blocking, device=idxs.device), [1, blocking]) + + # Duplicate for block row dimension. + idxs = torch.reshape(idxs, [idxs.size()[0], 1, blocking, 2]) + idxs = idxs.repeat(1, blocking, 1, 1) + + # Update the row indices. + idxs[:, :, :, 0] *= blocking + idxs[:, :, :, 0] += torch.reshape(torch.arange(blocking, device=idxs.device), [1, blocking, 1]) + idxs = torch.reshape(idxs, [-1, 2]) + return idxs + + +# TODO(tgale): Add input type checking. +@torch.no_grad() +def to_dense(x): + assert isinstance(x, Matrix) + + shape = (np.prod(x.shape[:-1]), x.shape[-1]) + row_idxs = x.row_indices.type(torch.int32) + col_idxs = x.column_indices.type(torch.int32) + indices = _expand_for_blocking(torch.stack([row_idxs, col_idxs], dim=1), x.blocking) + indices = (indices[:, 0] * shape[1] + indices[:, 1]).type(torch.int64) + + out = torch.zeros(shape[0] * shape[1], dtype=x.dtype, device=x.device) + out.scatter_(0, indices, x.data.flatten()) + return out.reshape(x.size()) + + +@torch.no_grad() +def _mask(x, blocking=1): + assert x.dim() == 2 + assert x.size()[0] % blocking == 0 + assert x.size()[1] % blocking == 0 + block_rows = x.size()[0] // blocking + block_cols = x.size()[1] // blocking + x = torch.reshape(x, [block_rows, blocking, block_cols, blocking]) + x = torch.sum(torch.abs(x), dim=(1, 3)) + return x != 0 + + +# TODO(tgale): Add input type checking. +@torch.no_grad() +def to_sparse(x, blocking=1): + m = _mask(x, blocking) + + # TODO(tgale): Set to appropriate type for input matrix. + row_nnzs = torch.sum(m, dim=1).type(torch.int32) + zeros = torch.zeros((1,), dtype=row_nnzs.dtype, device=row_nnzs.device) + offsets = torch.cat([zeros, torch.cumsum(row_nnzs, dim=0)]) + offsets = offsets.type(torch.int32) + + indices = torch.nonzero(m).type(torch.int16) + row_indices = indices[:, 0] + column_indices = indices[:, 1] + + # Nonzero indices in the dense matrix. + nonzero_indices = torch.nonzero(m) + nonzero_indices = _expand_for_blocking(nonzero_indices, blocking) + nonzero_indices = nonzero_indices[:, 0] * x.size()[1] + nonzero_indices[:, 1] + + # Gather the data and construct the sparse matrix. + data = torch.gather(x.flatten(), dim=0, index=nonzero_indices) + data = torch.reshape(data, [-1, blocking, blocking]) + return Matrix(x.size(), data, row_indices, column_indices, offsets) + + +@torch.no_grad() +def ones_like(x): + return Matrix(x.size(), + torch.ones_like(x.data), + x.row_indices, + x.column_indices, x.offsets) + + +def sum(x): + assert isinstance(x, Matrix) + return x.data.sum() diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/matrix_ops_test.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/matrix_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..d172d921f6f08b0e4fb709207a458b0e1e071bd0 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/ops/matrix_ops_test.py @@ -0,0 +1,62 @@ +import unittest + +# from absl.testing import parameterized +import stk +import torch + + +# @parameterized.parameters( +# (8, 16, 0.0, 1), +# (8, 16, 0.5, 1), +# (8, 16, .95, 1), +# (16, 8, 0.0, 1), +# (16, 8, 0.5, 1), +# (16, 8, .95, 1), +# (8, 16, 0.0, 8), +# (8, 16, 0.5, 8), +# (8, 16, 1.0, 8), +# (16, 8, 0.0, 8), +# (16, 8, 0.5, 8), +# (16, 8, 1.0, 8), +# (128, 256, 0.5, 16), +# (256, 128, 0.75, 32), +# (512, 512, .875, 128)) +# class MatrixOpsTest(parameterized.TestCase): +# +# def testMatrixOps_FormatConversion(self, rows, cols, sparsity, blocking): +# mask = stk.random.dense_mask(rows, cols, sparsity, blocking) +# x = (torch.randn(rows, cols) * mask).type(torch.float16) +# +# # Convert the matrix to sparse format. +# sparse_x = stk.ops.to_sparse(x, blocking) +# +# # Validate the matrix. +# sparse_x.validate() +# +# # Validate the shape. +# self.assertEqual(sparse_x.dim(), 2) +# self.assertEqual(sparse_x.size()[0], rows) +# self.assertEqual(sparse_x.size()[1], cols) +# +# # Validate the sparsity. +# numblocks = rows // blocking * cols // blocking +# nnz = round(numblocks * (1 - sparsity)) * blocking ** 2 +# self.assertEqual(sparse_x.nnz, nnz) +# +# # Convert back to dense format. +# dense_x = stk.ops.to_dense(sparse_x) +# +# # Validate the shape. +# self.assertEqual(dense_x.dim(), 2) +# self.assertEqual(dense_x.size()[0], rows) +# self.assertEqual(dense_x.size()[1], cols) +# +# # Validate the sparsity +# self.assertEqual(torch.count_nonzero(dense_x).item(), nnz) +# +# # Validate the output. +# self.assertTrue(torch.all(torch.eq(x, dense_x))) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/random/__init__.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/random/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2576d1ca27283f77569a9a620c7c99fa68aaf30e --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/random/__init__.py @@ -0,0 +1,2 @@ +# from stk.random.random_ops import dense_mask, mask, randn +from .random_ops import dense_mask, mask, randn diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/random/random_ops.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/random/random_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..d1b36771e0eb8e7abf46bcb3b136b5fb1d29df93 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/random/random_ops.py @@ -0,0 +1,36 @@ +import numpy as np +import torch +from ..ops import matrix_ops + + +@torch.no_grad() +def dense_mask(rows, cols, sparsity, blocking=1): + assert sparsity >= 0.0 and sparsity <= 1.0 + assert rows % blocking == 0 and cols % blocking == 0 + + block_rows, block_cols = (rows // blocking, cols // blocking) + nnz = round(block_rows * block_cols * (1 - sparsity)) + + out = np.ones(block_rows * block_cols) + mask = np.random.choice(out.size, out.size - nnz, replace=False) + out[mask] = 0.0 + + out = np.tile( + np.reshape(out, [block_rows, 1, block_cols, 1]), + (1, blocking, 1, blocking)) + out = np.reshape(out, [rows, cols]) + return torch.from_numpy(out.astype(np.float32)) + + +@torch.no_grad() +def mask(m, n, sparsity, blocking=1): + out = dense_mask(m, n, sparsity, blocking).type(torch.float16) + return matrix_ops.to_sparse(out, blocking=blocking) + + +@torch.no_grad() +def randn(shape, sparsity, blocking=1): + shape_2d = (np.prod(shape[:-1]), shape[-1]) + out = mask(*shape_2d, sparsity, blocking) + out.data.copy_(torch.randn(*out.data.shape)) + return out.view(*shape) diff --git a/build/torch214-cxx11-cu130-x86_64-linux/stk/random/random_ops_test.py b/build/torch214-cxx11-cu130-x86_64-linux/stk/random/random_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..71d716b78b5ec009cbf9ac2dfdf09162a0102e62 --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/stk/random/random_ops_test.py @@ -0,0 +1,73 @@ +import unittest + +# from absl.testing import parameterized +from . import random +import torch + + +# @parameterized.parameters( +# (8, 16, 0.0, 1), +# (8, 16, 0.5, 1), +# (8, 16, .95, 1), +# (16, 8, 0.0, 1), +# (16, 8, 0.5, 1), +# (16, 8, .95, 1), +# (8, 16, 0.0, 8), +# (8, 16, 0.5, 8), +# (8, 16, 1.0, 8), +# (16, 8, 0.0, 8), +# (16, 8, 0.5, 8), +# (16, 8, 1.0, 8), +# (128, 256, 0.5, 16), +# (256, 128, 0.75, 32), +# (512, 512, .875, 128)) +# class RandomOpsTest(parameterized.TestCase): +# +# def testRandomOps_DenseMask(self, rows, cols, sparsity, blocking): +# mask = random.dense_mask( +# rows, cols, sparsity, blocking) +# +# # Validate the shape. +# self.assertEqual(mask.dim(), 2) +# self.assertEqual(mask.size()[0], rows) +# self.assertEqual(mask.size()[1], cols) +# +# # Validate the sparsity +# numblocks = rows // blocking * cols // blocking +# nnz = round(numblocks * (1 - sparsity)) * blocking ** 2 +# self.assertEqual( +# torch.count_nonzero(mask).item(), +# nnz) +# +# # Check values are zero or one. +# self.assertTrue( +# torch.all(torch.logical_or( +# torch.eq(mask, 0), +# torch.eq(mask, 1)))) +# +# def testRandomOps_SparseMask(self, rows, cols, sparsity, blocking): +# mask = random.mask( +# rows, cols, sparsity, blocking) +# +# # Validate the matrix. +# mask.validate() +# +# # Validate the shape. +# self.assertEqual(mask.dim(), 2) +# self.assertEqual(mask.size()[0], rows) +# self.assertEqual(mask.size()[1], cols) +# +# # Validate the sparsity. +# numblocks = rows // blocking * cols // blocking +# nnz = round(numblocks * (1 - sparsity)) * blocking ** 2 +# self.assertEqual(mask.nnz, nnz) +# +# # Check values are zero or one. +# self.assertTrue( +# torch.all(torch.logical_or( +# torch.eq(mask.data, 0), +# torch.eq(mask.data, 1)))) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu130-x86_64-linux/xpu_fused_moe.py b/build/torch214-cxx11-cu130-x86_64-linux/xpu_fused_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..b709e394e2e0461d858dec0f58176c0543a15e3c --- /dev/null +++ b/build/torch214-cxx11-cu130-x86_64-linux/xpu_fused_moe.py @@ -0,0 +1,717 @@ +# SPDX-License-Identifier: Apache-2.0 +# MegaBlocks XPU Fused MoE Implementation +import os +import torch + +from ._ops import ops, add_op_namespace_prefix + +from torch.library import register_fake + + +def resolve_dtensor(weight: torch.Tensor): + """Convert DTensor to local tensor for use with custom ops.""" + from torch.distributed._tensor import DTensor + if isinstance(weight, DTensor): + return weight.to_local() + return weight + + +# Register fake/meta kernels for torch.compile compatibility +def _register_xpu_fake_kernels(): + """Register fake kernels for XPU MoE operations to support torch.compile.""" + + def _register_if_available(op_name, fn): + if hasattr(ops, op_name): + register_fake(add_op_namespace_prefix(op_name))(fn) + + _register_if_available( + "cutlass_grouped_gemm_interface", + lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4, is_B_mxfp8=False: ptr_D, + ) + + _register_if_available( + "fused_moe_prologue", + lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, ep_rank, ep_size, num_experts_on_rank: None, + ) + + _register_if_available( + "moe_gather", + lambda output, moe_output, topk_weights, permuted_row_to_unpermuted_row, unpermuted_row_to_permuted_row, expert_first_token_offset, num_experts: None, + ) + + _register_if_available( + "silu_and_mul", + lambda out, input: None, + ) + _register_if_available( + "mul_and_silu", + lambda out, input: None, + ) + _register_if_available( + "gelu_and_mul", + lambda out, input: None, + ) + _register_if_available( + "gelu_tanh_and_mul", + lambda out, input: None, + ) + _register_if_available( + "gelu_fast", + lambda out, input: None, + ) + _register_if_available( + "gelu_new", + lambda out, input: None, + ) + _register_if_available( + "gelu_quick", + lambda out, input: None, + ) + _register_if_available( + "swigluoai_and_mul", + lambda out, input, alpha=1.702, limit=7.0: None, + ) + + +# Register fake kernels on module load +_register_xpu_fake_kernels() + + +# default +def cutlass_grouped_gemm(input_A, input_B, bias, output, expert_token_count, n, + k, num_experts): + # expert_token_count_ = torch.tensor(expert_token_count, + # dtype=torch.int64, + # device=input_A.device) + # if bias is not None: + # bias = bias.repeat_interleave(expert_token_count_, dim=0).float() + + def exclusive_prefix_sum(arr): + prefix = [0] + for i, x in enumerate(arr): + prefix.append(prefix[-1] + x) + return prefix + + expert_offset = torch.tensor(exclusive_prefix_sum(expert_token_count), + dtype=torch.int64, + device="xpu") + ops.cutlass_grouped_gemm_interface( + ptr_A=input_A, + ptr_B=input_B, + ptr_scales=None, + ptr_bias=bias, + ptr_D=output, + expert_first_token_offset=expert_offset, + N=n, + K=k, + num_experts=num_experts, + is_B_int4=False, + is_B_mxfp4=False, + is_B_mxfp8=False) + + +def cutlass_grouped_gemm_xe2(input_A, input_B, scales, bias, output, + num_rows_per_expert, n, k, num_experts, is_B_int4, + is_B_mxfp4, is_B_mxfp8=False): + expert_first_token_offset = torch.cat([ + torch.tensor([0], + dtype=num_rows_per_expert.dtype, + device=num_rows_per_expert.device), + torch.cumsum(num_rows_per_expert, dim=0) + ]).to(torch.int64) + ops.cutlass_grouped_gemm_interface( + ptr_A=input_A, + ptr_B=input_B, + ptr_scales=scales, + ptr_bias=bias, + ptr_D=output, + expert_first_token_offset=expert_first_token_offset, + N=n, + K=k, + num_experts=num_experts, + is_B_int4=is_B_int4, + is_B_mxfp4=is_B_mxfp4, + is_B_mxfp8=is_B_mxfp8) + + +def ceilDiv(a, b): + return (a + b - 1) // b + + +def compute_num_tokens_per_block(num_tokens, num_experts_per_node): + for num_tokens_per_block in [32, 64, 128, 256, 512, 1024]: + num_blocks_per_seq = ceilDiv(num_tokens, num_tokens_per_block) + if num_blocks_per_seq * num_experts_per_node <= num_tokens_per_block: + return num_tokens_per_block + return 1024 + + +def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """Reinterpret a uint8 buffer as a typed tensor by copying bytes. + + This avoids `Tensor.view(dtype)` which can fail under torch.compile + constant folding when shape divisibility is not proven. + """ + if byte_tensor.dtype != torch.uint8: + raise ValueError("byte_tensor must be uint8") + itemsize = torch.empty((), dtype=dtype).element_size() + numel = byte_tensor.numel() // itemsize + out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device) + out.view(torch.uint8).copy_(byte_tensor.contiguous()) + return out + + +def implement_zp(qweight): + # change u4 to s4 to avoid zero point in gemm kernel + # only support default zero point now + assert qweight.dtype == torch.uint8, "Input tensor must be uint8" + + high_u4 = (qweight >> 4) & 0x0F + low_u4 = qweight & 0x0F + + high_s8 = high_u4.to(torch.int8) + low_s8 = low_u4.to(torch.int8) + + high_s8 = high_s8 - 8 + low_s8 = low_s8 - 8 + + def pack_compact(a, b): + + def process_number(x): + sign = (x < 0).to(torch.uint8) + abs_low3 = (x.view(torch.uint8) & 0x7).to(torch.uint8) + return (sign << 3) | abs_low3 + + packed_a = process_number(a) + packed_b = process_number(b) + + return (packed_a << 4) | packed_b + + result = pack_compact(high_s8, low_s8) + + return result + + +def xpu_fused_moe(hidden_states, + w13, + w13_scales, + w13_bias, + w2, + w2_scales, + w2_bias, + topk_weights, + topk_ids, + n_experts_per_token, + activation, + num_experts, + ep_rank=0, + ep_size=1, + is_fp8=False, + is_int4=False, + is_mxfp4=False, + is_mxfp8=False): + ''' + hidden_states: [num_rows, hidden_size] + w13: [num_experts, 2*inter_size, hidden_size] + w13_scales: + None for bf16/fp16 + or [num_experts] for fp8 + or [num_experts, 2*inter_size, hidden_size // group_size] for 4bits + w13_bias: [num_experts, 2*inter_size] or None + w2: [num_experts, hidden_size, inter_size] + w2_scales: + None for bf16/fp16 + or [num_experts] for fp8 + or [num_experts, hidden_size, inter_size // group_size] for 4bits + w2_bias: [num_experts, hidden_size] or None + topk_weights: [num_rows, topk] + topk_ids: [num_rows, topk] + n_experts_per_token: int + activation: str + num_experts: int + is_int4: bool + is_mxfp4: bool + is_mxfp8: bool + ''' + + # Resolve DTensors to local tensors before passing to custom ops + hidden_states = resolve_dtensor(hidden_states) + w13 = resolve_dtensor(w13) + w2 = resolve_dtensor(w2) + if w13_scales is not None: + w13_scales = resolve_dtensor(w13_scales) + if w13_bias is not None: + w13_bias = resolve_dtensor(w13_bias) + if w2_scales is not None: + w2_scales = resolve_dtensor(w2_scales) + if w2_bias is not None: + w2_bias = resolve_dtensor(w2_bias) + topk_weights = resolve_dtensor(topk_weights) + topk_ids = resolve_dtensor(topk_ids) + + output = torch.empty_like(hidden_states) + num_rows, hidden_size = list(hidden_states.shape) + + dim_last = w13.shape[-1] + dim_second_last = w13.shape[-2] + + # w13 is combined gate+up weights, so one dimension is 2*inter_size + # Determine which dimension is hidden_size and which is 2*inter_size + if dim_second_last == hidden_size: + # w13 is [E, hidden_size, 2*inter_size] - standard layout + inter_size = dim_last // 2 + needs_transpose = False + else: + # w13 is [E, 2*inter_size, hidden_size] - needs transpose + inter_size = dim_second_last // 2 + needs_transpose = True + + assert w13.is_contiguous() and w2.is_contiguous() + + # 4bits support [E, N, K] + # other types [E, K, N] + if not is_int4 and not is_mxfp4 and not is_mxfp8: + if not hasattr(w13, 'xpu_fused_moe'): + if needs_transpose: + w13.data = w13.transpose(-1, -2).contiguous() + w2.data = w2.transpose(-1, -2).contiguous() + w13.xpu_fused_moe = True + w13.inter_size = inter_size + else: + inter_size = w13.inter_size + + if is_int4 and not hasattr(w13, 'xpu_fused_moe'): + w13_tmp = torch.empty_like(w13) + w2_tmp = torch.empty_like(w2) + for i in range(num_experts): + w13_tmp[i] = implement_zp(w13[i]) + w2_tmp[i] = implement_zp(w2[i]) + w13_tmp = w13_tmp.contiguous() + w2_tmp = w2_tmp.contiguous() + w13.data = w13_tmp + w2.data = w2_tmp + w13.xpu_fused_moe = True + + # TODO: will all integrated in Cpp func. Temporary expose before gemm fusion + num_experts_per_node = num_experts + experts_per_token = n_experts_per_token + num_moe_inputs = n_experts_per_token * num_rows + permuted_elems = num_moe_inputs * hidden_size + # interbuf_elems = num_moe_inputs * inter_size + permuted_row_to_unpermuted_row_size = num_moe_inputs * 4 + permuted_token_selected_experts_size = num_moe_inputs * 4 + src_to_dest_map_size = experts_per_token * num_rows * 4 + expert_first_token_offset_size = (num_experts_per_node + 1) * 8 + num_tokens_per_block = compute_num_tokens_per_block( + num_rows, num_experts_per_node) + num_blocks_per_seq = ceilDiv(num_rows, num_tokens_per_block) + blocked_expert_counts_size = num_experts_per_node * num_blocks_per_seq * 4 + blocked_expert_counts_cumsum_size = blocked_expert_counts_size + blocked_row_to_unpermuted_row_size = num_experts_per_node * num_rows * 4 + permuted_data_size = permuted_elems * hidden_states.element_size() + permuted_token_final_scales_size = num_moe_inputs * 4 + + ws_map = {} + map_offset = 0 + + def config_ws(name, size): + nonlocal map_offset + if size % 256 != 0: + size += 256 - size % 256 + ws_map[name] = (size, map_offset) + map_offset += size + + config_ws("permuted_row_to_unpermuted_row", + permuted_row_to_unpermuted_row_size) + config_ws("permuted_token_selected_experts", + permuted_token_selected_experts_size) + config_ws("unpermuted_row_to_permuted_row", src_to_dest_map_size) + config_ws("blocked_expert_counts", blocked_expert_counts_size) + config_ws("blocked_expert_counts_cumsum", + blocked_expert_counts_cumsum_size) + config_ws("blocked_row_to_unpermuted_row", + blocked_row_to_unpermuted_row_size) + config_ws("expert_first_token_offset", expert_first_token_offset_size) + config_ws("permuted_token_final_scales", permuted_token_final_scales_size) + config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size) + + workspace = torch.zeros(map_offset, + dtype=torch.uint8, + device=hidden_states.device) + if topk_ids.dtype == torch.int32: + topk_ids = topk_ids.to(torch.int64) + ops.fused_moe_prologue( + input=hidden_states, + token_selected_experts=topk_ids, + token_final_scales=topk_weights, + workspace=workspace, + hidden_size=hidden_size, + inter_size=inter_size, + ep_rank=ep_rank, + ep_size=ep_size, + num_experts_on_rank=num_experts_per_node) + + expert_first_token_offset_bytes = workspace[ + ws_map["expert_first_token_offset"][1]: + ws_map["expert_first_token_offset"][1] + + expert_first_token_offset_size] + unpermuted_row_to_permuted_row_bytes = workspace[ + ws_map["unpermuted_row_to_permuted_row"][1]: + ws_map["unpermuted_row_to_permuted_row"][1] + + src_to_dest_map_size] + permuted_row_to_unpermuted_row_bytes = workspace[ + ws_map["permuted_row_to_unpermuted_row"][1]: + ws_map["permuted_row_to_unpermuted_row"][1] + + permuted_row_to_unpermuted_row_size] + + if torch.compiler.is_compiling(): + expert_first_token_offset = _bytes_to_typed_tensor( + expert_first_token_offset_bytes, torch.int64 + ) + unpermuted_row_to_permuted_row = _bytes_to_typed_tensor( + unpermuted_row_to_permuted_row_bytes, torch.int32 + ) + permuted_row_to_unpermuted_row = _bytes_to_typed_tensor( + permuted_row_to_unpermuted_row_bytes, torch.int32 + ) + else: + expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64) + unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32) + permuted_row_to_unpermuted_row = permuted_row_to_unpermuted_row_bytes.view(torch.int32) + gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]: + ws_map["overlapped_gemm1_gemm2_inputs"][1] + + permuted_data_size].view(hidden_states.dtype).view( + num_moe_inputs, hidden_size) + # permuted_token_final_scales = workspace[ + # ws_map["permuted_token_final_scales"][1]: + # ws_map["permuted_token_final_scales"][1] + + # permuted_token_final_scales_size].view(torch.float) + gemm1_output = torch.empty((num_moe_inputs, 2 * inter_size), + dtype=hidden_states.dtype, + device=hidden_states.device) + + ########### gemm1 ################## + input_B = w13 + + if not is_fp8 and not is_int4 and not is_mxfp4 and not is_mxfp8: + ops.cutlass_grouped_gemm_interface( + ptr_A=gemm1_input, + ptr_B=input_B, + ptr_scales=None, + ptr_bias=w13_bias, + ptr_D=gemm1_output, + expert_first_token_offset=expert_first_token_offset, + N=2 * inter_size, + K=hidden_size, + num_experts=num_experts_per_node, + is_B_int4=is_int4, + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) + else: + ops.cutlass_grouped_gemm_interface( + ptr_A=gemm1_input, + ptr_B=input_B, + ptr_scales=w13_scales, + ptr_bias=w13_bias, + ptr_D=gemm1_output, + expert_first_token_offset=expert_first_token_offset, + N=2 * inter_size, + K=hidden_size, + num_experts=num_experts_per_node, + is_B_int4=is_int4, + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) + + # act + act_output = torch.empty((num_moe_inputs, inter_size), + dtype=gemm1_output.dtype, + device=gemm1_output.device) + if activation == "silu": + ops.silu_and_mul(act_output, gemm1_output) + elif activation == "gelu": + ops.gelu_and_mul(act_output, gemm1_output) + elif activation == "swigluoai": + ops.swigluoai_and_mul(act_output, gemm1_output, 1.702, 7.0) + else: + raise ValueError(f"Unsupported FusedMoe activation: {activation}.") + + ########### gemm2 ################## + input_A = act_output.contiguous() + input_B = w2 + gemm2_output = torch.empty((num_moe_inputs, hidden_size), + dtype=hidden_states.dtype, + device=hidden_states.device) + if not is_fp8 and not is_int4 and not is_mxfp4 and not is_mxfp8: + ops.cutlass_grouped_gemm_interface( + ptr_A=input_A, + ptr_B=input_B, + ptr_scales=None, + ptr_bias=w2_bias, + ptr_D=gemm2_output, + expert_first_token_offset=expert_first_token_offset, + N=hidden_size, + K=inter_size, + num_experts=num_experts_per_node, + is_B_int4=is_int4, + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) + else: + ops.cutlass_grouped_gemm_interface( + ptr_A=input_A, + ptr_B=input_B, + ptr_scales=w2_scales, + ptr_bias=w2_bias, + ptr_D=gemm2_output, + expert_first_token_offset=expert_first_token_offset, + N=hidden_size, + K=inter_size, + num_experts=num_experts_per_node, + is_B_int4=is_int4, + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) + + ops.moe_gather(output, gemm2_output, topk_weights, + permuted_row_to_unpermuted_row, + unpermuted_row_to_permuted_row, + expert_first_token_offset, + num_experts_per_node) + return output + + +def apply_jitter(x: torch.Tensor, moe_jitter_eps: float) -> torch.Tensor: + """Apply jitter to the input tensor for regularization.""" + low = 1.0 - moe_jitter_eps + high = 1.0 + moe_jitter_eps + noise = torch.rand(x.size(), dtype=x.dtype, device=x.device) + return x * (low + noise * (high - low)) + + +def compute_top_k(scores: torch.Tensor, moe_top_k: int): + """Compute the top-k scores from the logits.""" + if moe_top_k == 1: + return scores.max(dim=-1, keepdim=True) + return torch.topk(scores, moe_top_k, dim=-1) + + +def route_tokens_xpu( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + moe_top_k: int, + moe_num_experts: int, + moe_jitter_eps: float = None, + moe_normalize_expert_weights: int = None, + training: bool = False, +) -> tuple: + """Route tokens to experts and compute expert weights and indices (XPU version).""" + if training and moe_jitter_eps is not None: + x = apply_jitter(x, moe_jitter_eps) + + x_flat = x.view(-1, x.shape[-1]) + logits = torch.nn.functional.linear(x_flat, router_weight, router_bias) + expert_weights, expert_indices = compute_top_k(logits, moe_top_k) + expert_weights = expert_weights.softmax(dim=-1) + if moe_normalize_expert_weights is not None: + expert_weights = expert_weights / torch.norm( + expert_weights, + p=moe_normalize_expert_weights, + dim=-1, + keepdim=True, + ) + + return logits, expert_weights, expert_indices + + +def _get_device_mesh(model): + """Extract device_mesh from child's unused pre_hook closure for EP support.""" + try: + hook = next( + h + for h in model.experts._forward_pre_hooks.values() + if "device_mesh" in h.__code__.co_freevars + ) + return hook.__closure__[ + hook.__code__.co_freevars.index("device_mesh") + ].cell_contents + except Exception: + return None + + +def _mxfp4_expert_tensors(experts, dtype): + """View transformers' MXFP4 experts as the tensors the SYCL grouped GEMM expects. + + Weights and scales arrive wrapped in triton tensors. On XPU both use StridedLayout, + whose swizzle is the identity, so `storage.data` is still the checkpoint tensor - just + transposed to [E, K/2, N]. Transposing back therefore restores the original contiguous + [E, N, K/2] weights and [E, N, K/32] E8M0 scales as plain views: `contiguous()` finds + them already contiguous and returns them untouched, so this costs no copy and no extra + memory. The experts module is left alone, which keeps the triton tensors and the + PrecisionConfig that `save_pretrained` needs to re-serialize the checkpoint. + """ + tensors = [] + for proj in ("gate_up_proj", "down_proj"): + config = getattr(experts, f"{proj}_precision_config") + tensors.append(getattr(experts, proj).storage.data.transpose(-1, -2).contiguous()) + tensors.append(config.weight_scale.storage.data.transpose(-1, -2).contiguous()) + + # The GEMM reads the bias as ElementA; `to` is a no-op when it already matches. + bias = getattr(experts, f"{proj}_bias", None) + tensors.append(None if bias is None else bias.data.to(dtype)) + + return tensors + + +class MegaBlocksMoeMLP(torch.nn.Module): + can_torch_compile: bool = True + + def forward(self, x: torch.Tensor) -> tuple: + """ + Forward pass through the MoE layer. + + Args: + x: Input tensor of shape [batch_size, seq_len, hidden_size] or [tokens, hidden_size] + + Returns: + Tuple of (output, expert_weights) where: + - output: Tensor of same shape as input + - expert_weights: Expert weights for each token [tokens, top_k] + """ + + # Get MoE parameters from the wrapped modules + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + moe_jitter_eps = getattr(self.experts, "jitter_eps", None) + moe_normalize_expert_weights = getattr( + self.experts, "normalize_expert_weights", None + ) + + # Get EP (Expert Parallelism) parameters + ep_size = 1 + ep_rank = 0 + expert_parallel_group = getattr(self, "expert_parallel_group", None) + if expert_parallel_group is None: + device_mesh = _get_device_mesh(self) + if device_mesh is not None: + expert_parallel_group = device_mesh.get_group() + if expert_parallel_group is not None: + import torch.distributed as dist + if dist.is_initialized(): + ep_size = dist.get_world_size(expert_parallel_group) + ep_rank = dist.get_rank(expert_parallel_group) + + # Number of experts on this rank + num_experts_on_rank = moe_num_experts // ep_size + + # Detect activation type - check for GptOss-style swigluoai activation + # GptOssExperts has alpha and limit attributes for swigluoai + if hasattr(self.experts, "alpha") and hasattr(self.experts, "limit"): + activation = "swigluoai" + else: + activation = getattr(self.experts, "activation", "silu") + + # Get weight tensors - support different naming conventions + if hasattr(self.experts, "gate_up_proj"): + w13 = self.experts.gate_up_proj + # NOTE: swigluoai_and_mul kernel expects interleaved layout [g0,u0,g1,u1,...] + # which matches GptOss's gate_up_proj layout, so no conversion needed. + + elif hasattr(self.experts, "w1"): + # Combine w1 and w3 if stored separately + w1 = self.experts.w1 + w3 = getattr(self.experts, "w3", None) + if w3 is not None: + w13 = torch.cat([w1, w3], dim=-2) + else: + w13 = w1 + else: + raise AttributeError("experts module must have 'gate_up_proj' or 'w1' attribute") + + if hasattr(self.experts, "down_proj"): + w2 = self.experts.down_proj + elif hasattr(self.experts, "w2"): + w2 = self.experts.w2 + else: + raise AttributeError("experts module must have 'down_proj' or 'w2' attribute") + + # Get optional bias tensors + w13_bias = getattr(self.experts, "gate_up_proj_bias", None) + w2_bias = getattr(self.experts, "down_proj_bias", None) + + # Get quantization info + is_fp8 = getattr(self.experts, "is_fp8", False) + is_int4 = getattr(self.experts, "is_int4", False) + is_mxfp4 = getattr(self.experts, "is_mxfp4", False) + is_mxfp8 = getattr(self.experts, "is_mxfp8", False) + + w13_scales = getattr(self.experts, "gate_up_proj_scales", None) + w2_scales = getattr(self.experts, "down_proj_scales", None) + + # Native MXFP4 checkpoint: derive the GEMM's view of the experts once and cache it. + mxfp4 = getattr(self, "_mxfp4_tensors", None) + if (mxfp4 is None or mxfp4[0] is not x.dtype) and getattr( + self.experts, "gate_up_proj_precision_config", None + ) is not None: + mxfp4 = (x.dtype, *_mxfp4_expert_tensors(self.experts, x.dtype)) + self._mxfp4_tensors = mxfp4 + if mxfp4 is not None: + _, w13, w13_scales, w13_bias, w2, w2_scales, w2_bias = mxfp4 + is_mxfp4 = True + + # Store original shape + in_shape = x.size() + + # Route tokens to experts + logits, expert_weights, expert_indices = route_tokens_xpu( + x, + self.router.weight, + getattr(self.router, "bias", None), + moe_top_k, + moe_num_experts, + moe_jitter_eps, + moe_normalize_expert_weights, + self.training, + ) + + # Reshape input for fused MoE + x_flat = x.view(-1, x.shape[-1]) + + # Call XPU fused MoE kernel + output = xpu_fused_moe( + hidden_states=x_flat, + w13=w13, + w13_scales=w13_scales, + w13_bias=w13_bias, + w2=w2, + w2_scales=w2_scales, + w2_bias=w2_bias, + topk_weights=expert_weights.float(), + topk_ids=expert_indices, + n_experts_per_token=moe_top_k, + activation=activation, + num_experts=num_experts_on_rank, + ep_rank=ep_rank, + ep_size=ep_size, + is_fp8=is_fp8, + is_int4=is_int4, + is_mxfp4=is_mxfp4, + is_mxfp8=is_mxfp8, + ) + + # All-reduce across EP group to combine partial expert outputs + if ep_size > 1 and expert_parallel_group is not None: + import torch.distributed as dist + dist.all_reduce(output, op=dist.ReduceOp.SUM, group=expert_parallel_group) + + # Restore original shape + output = output.view(in_shape) + + return output, expert_weights + + +# Export classes and functions +__all__ = [ + "MegaBlocksMoeMLP", + "xpu_fused_moe", + "cutlass_grouped_gemm", + "cutlass_grouped_gemm_xe2", +] \ No newline at end of file diff --git a/build/torch214-cxx11-cu132-x86_64-linux/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0766d7b8da4f97baca212177b4bb989bc6374bf8 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/__init__.py @@ -0,0 +1,205 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch + +# Stable alias: bare `ops` is shadowed by `from . import layers` below. +from ._ops import ops as _compiled_ops +from . import ops + +from .grouped_gemm import backend as gg_backend +from .grouped_gemm import ops as gg_ops + + +from ._layers.arguments import Arguments +from ._layers.dmoe import ParallelDroplessMLP, dMoE +from ._layers.glu import SparseGLU +from ._layers.mlp import MLP, SparseMLP +from ._layers.moe import MoE, ParallelMLP, get_load_balancing_loss + +from . import layers + +# This section contains the direct kernel exports (not inlcuded in the original code) +def exclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor: + """ + Compute exclusive cumulative sum along the specified dimension. + + Args: + x: Input tensor + dim: Dimension along which to compute cumsum + out: Output tensor (modified in-place) + + Returns: + The output tensor + """ + result = ops.exclusive_cumsum(x, dim) + out.copy_(result) + return out + + +def inclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor: + """ + Compute inclusive cumulative sum along the specified dimension. + + Args: + x: Input tensor + dim: Dimension along which to compute cumsum + out: Output tensor (modified in-place) + + Returns: + The output tensor + """ + result = ops.inclusive_cumsum(x, dim) + out.copy_(result) + return out + + +def histogram(x: torch.Tensor, num_bins: int) -> torch.Tensor: + """ + Compute histogram of input tensor values. + + Args: + x: Input tensor + num_bins: Number of histogram bins + + Returns: + Histogram tensor with counts for each bin + """ + return ops.histogram(x, num_bins) + + +def indices( + padded_bins: torch.Tensor, + block_size: int, + output_block_rows: int, + output_block_columns: int, +) -> torch.Tensor: + """ + Construct indices from padded bins for sparse operations. + + Args: + padded_bins: Tensor containing bin boundaries + block_size: Size of each block + output_block_rows: Number of rows in output blocks + output_block_columns: Number of columns in output blocks + + Returns: + Tensor containing constructed indices + """ + return ops.indices(padded_bins, block_size, output_block_rows, output_block_columns) + + +def replicate_forward( + x: torch.Tensor, bins: torch.Tensor, out: torch.Tensor +) -> torch.Tensor: + """ + Forward pass of replicate operation - replicate values according to bin sizes. + + Args: + x: Input tensor with values to replicate + bins: Tensor containing bin sizes + out: Output tensor (modified in-place) + + Returns: + The output tensor + """ + return ops.replicate_forward(x, bins, out) + + +def replicate_backward( + grad: torch.Tensor, bins: torch.Tensor, out: torch.Tensor +) -> torch.Tensor: + """ + Backward pass of replicate operation - reduce gradients back to bins. + + Args: + grad: Gradient tensor to reduce + bins: Tensor containing bin sizes + out: Output tensor (modified in-place) + + Returns: + The output tensor + """ + return ops.replicate_backward(grad, bins, out) + + +def sort( + x: torch.Tensor, end_bit: int, x_out: torch.Tensor, iota_out: torch.Tensor +) -> torch.Tensor: + """ + Radix sort with index tracking. + + Args: + x: Input tensor to sort + end_bit: Number of bits to consider in sorting + x_out: Output tensor for sorted values + iota_out: Output tensor for sorted indices + + Returns: + The sorted values tensor + """ + _compiled_ops.sort(x, end_bit, x_out, iota_out) + return x_out + + +# Convenience functions for common use cases +def cumsum(x: torch.Tensor, dim: int = -1, exclusive: bool = False) -> torch.Tensor: + """ + Compute cumulative sum with automatic output allocation. + + Args: + x: Input tensor + dim: Dimension along which to compute cumsum (default: last dimension) + exclusive: Whether to compute exclusive (True) or inclusive (False) cumsum + + Returns: + New tensor containing the cumulative sum + """ + out = torch.empty_like(x) + if exclusive: + return exclusive_cumsum(x, dim, out) + else: + return inclusive_cumsum(x, dim, out) + + +def argsort(x: torch.Tensor, end_bit: int = 32) -> tuple[torch.Tensor, torch.Tensor]: + """ + Sort tensor and return both sorted values and indices. + + Args: + x: Input tensor to sort + end_bit: Number of bits to consider in sorting + + Returns: + Tuple of (sorted_values, sorted_indices) + """ + x_out = torch.empty_like(x) + iota_out = torch.empty_like(x) + sort(x, end_bit, x_out, iota_out) + return x_out, iota_out + + +# Export public API +__all__ = [ + "MyReplacementLayer", + # Direct kernel exports + "exclusive_cumsum", + "inclusive_cumsum", + "histogram", + "indices", + "replicate_forward", + "replicate_backward", + "sort", + "cumsum", + "argsort", + # Original exports + "Arguments", + "ParallelDroplessMLP", + "dMoE", + "SparseGLU", + "MLP", + "SparseMLP", + "MoE", + "ParallelMLP", + "get_load_balancing_loss", +] diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/_triton_kernels/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/_triton_kernels/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/_triton_kernels/gmm.py b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/_triton_kernels/gmm.py new file mode 100644 index 0000000000000000000000000000000000000000..8c101d07cea416f9390b708e5a35fdc466e48aed --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/_triton_kernels/gmm.py @@ -0,0 +1,574 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. + + +# Imports. +# ------------------------------------------------------------------------------ + +# Python standard library +import functools + +# Triton +import triton +import triton.language as tl + +# AITER +from ..configs import CONFIGS as _CONFIGS +from ..utils._triton import arch_info +from ..utils._triton.pid_preprocessing import pid_grid, remap_xcd + +# Kernel config. +# ------------------------------------------------------------------------------ + + +@functools.lru_cache() +def get_config( + gmm_type: str, M: int, K: int, N: int, G: int, accumulate: bool = False +) -> dict[str, int]: + assert gmm_type in { + "gmm", + "ptgmm", + "nptgmm", + }, f"'{gmm_type}' is an invalid GMM variant." + dev = arch_info.get_arch() + assert ( + dev in _CONFIGS + ), f"No GMM configuration tuned for arch '{dev}'. Supported: {sorted(_CONFIGS)}." + arch_configs = _CONFIGS[dev] + assert ( + "default" in arch_configs[gmm_type] + ), "Default configuration is absent." + key = "accumulate" if accumulate else "default" + return arch_configs[gmm_type][key] + + +# Common code shared by GMM and TGMM kernels. +# ------------------------------------------------------------------------------ + + +# XCD remapping followed by 1D PID to 2D grid mapping. +@triton.jit +def _remap_xcd_tile_grid( + tile_in_mm, + num_row_tiles, + num_col_tiles, + GROUP_SIZE: tl.constexpr = 1, + NUM_XCDS: tl.constexpr = 8, +): + return pid_grid( + remap_xcd(tile_in_mm, num_row_tiles * num_col_tiles, NUM_XCDS=NUM_XCDS), + num_row_tiles, + num_col_tiles, + GROUP_SIZE_M=GROUP_SIZE, + ) + + +# GMM kernel. +# ------------------------------------------------------------------------------ + + +@triton.heuristics( + { + "K_DIVISIBLE_BY_BLOCK_SIZE_K": lambda META: META["K"] % META["BLOCK_SIZE_K"] + == 0, + } +) +@triton.jit +def gmm_kernel( + # Tensor pointers: + lhs_ptr, + rhs_ptr, + group_sizes_ptr, + out_ptr, + bias_ptr, + # Tensor shapes: + M: int, + K: int, + N: int, + G: int, + # Meta-parameters: + TRANS_RHS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + K_DIVISIBLE_BY_BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE: tl.constexpr, + GRID_DIM: tl.constexpr, + USE_BIAS: tl.constexpr, +): + tl.assume(M > 0) + tl.assume(K > 0) + tl.assume(N > 0) + tl.assume(G > 0) + + num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) + tl.device_assert(num_n_tiles > 0, "num_n_tiles <= 0") + + # Current tile. Each program computes multiple tiles of each group. + tile = tl.program_id(0) + tl.device_assert(tile >= 0, "tile < 0 (at initialization)") + + # Tile limit of last MM problem (inclusive). + last_mm_tile = 0 + + # Last input row of lhs and output row of out. Each group reads some rows of + # lhs and writes some rows to out. + last_m = 0 + + # Loop through all (m, K, N) MM problems: + # (m, K) x (K, N) = (m, N) + # sum(m) = M + for g in range(G): + # Get m dimension of current MM problem. + m = tl.load(group_sizes_ptr + g) + # m can be zero if group is empty + tl.device_assert(m >= 0, "m < 0") + + num_m_tiles = tl.cdiv(m, BLOCK_SIZE_M) + # num_m_tiles can be zero if group is empty + tl.device_assert(num_m_tiles >= 0, "num_m_tiles < 0") + + num_tiles = num_m_tiles * num_n_tiles + # num_tiles can be zero if group is empty + tl.device_assert(num_tiles >= 0, "num_tiles < 0") + + # Loop through tiles of current MM problem. + while tile >= last_mm_tile and tile < last_mm_tile + num_tiles: + # Figure out tile coordinates in current MM problem. + tile_in_mm = tile - last_mm_tile + tl.device_assert(tile_in_mm >= 0, "tile_in_mm < 0") + + tile_m, tile_n = _remap_xcd_tile_grid( + tile_in_mm, num_m_tiles, num_n_tiles, GROUP_SIZE=GROUP_SIZE + ) + + # Do regular MM: + + tl.device_assert(tile_m * BLOCK_SIZE_M >= 0, "tile_m * BLOCK_SIZE_M < 0") + tl.device_assert(tile_n * BLOCK_SIZE_N >= 0, "tile_n * BLOCK_SIZE_N < 0") + + offs_lhs_m = ( + tile_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + ) % m + offs_rhs_n = ( + tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + ) % N + offs_k = tl.arange(0, BLOCK_SIZE_K).to(tl.int64) + + lhs_ptrs = lhs_ptr + (last_m + offs_lhs_m[:, None]) * K + offs_k[None, :] + + if TRANS_RHS: + rhs_ptrs = ( + rhs_ptr + + g.to(tl.int64) * K * N + + offs_k[:, None] + + offs_rhs_n[None, :] * K + ) + else: + rhs_ptrs = ( + rhs_ptr + + g.to(tl.int64) * K * N + + offs_k[:, None] * N + + offs_rhs_n[None, :] + ) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + if K_DIVISIBLE_BY_BLOCK_SIZE_K: + lhs = tl.load(lhs_ptrs) + rhs = tl.load(rhs_ptrs) + else: + k_mask_limit = K - k * BLOCK_SIZE_K + lhs = tl.load( + lhs_ptrs, mask=offs_k[None, :] < k_mask_limit, other=0 + ) + rhs = tl.load( + rhs_ptrs, mask=offs_k[:, None] < k_mask_limit, other=0 + ) + + acc = tl.dot(lhs, rhs, acc=acc) + + lhs_ptrs += BLOCK_SIZE_K + + if TRANS_RHS: + rhs_ptrs += BLOCK_SIZE_K + else: + rhs_ptrs += BLOCK_SIZE_K * N + + # Add bias if enabled + if USE_BIAS: + offs_bias_n = tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange( + 0, BLOCK_SIZE_N + ) + bias_ptrs = bias_ptr + g.to(tl.int64) * N + offs_bias_n + bias = tl.load(bias_ptrs, mask=offs_bias_n < N, other=0.0) + # Convert bias to float32 to match accumulator precision + bias = bias.to(tl.float32) + # Broadcast bias across M dimension and add in float32 + acc += bias[None, :] + + # Convert to output dtype after all computations + acc = acc.to(out_ptr.type.element_ty) + + offs_out_m = tile_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_out_n = tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + + out_ptrs = ( + out_ptr + (last_m + offs_out_m[:, None]) * N + offs_out_n[None, :] + ) + + tl.store( + out_ptrs, + acc, + mask=(offs_out_m[:, None] < m) & (offs_out_n[None, :] < N), + ) + + # Go to the next tile by advancing number of programs. + tile += GRID_DIM + tl.device_assert(tile > 0, "tile <= 0 (at update)") + + # Get ready to go to the next MM problem. + + last_mm_tile += num_tiles + # last_mm_tile can be zero if group 0 is skipped + tl.device_assert(last_mm_tile >= 0, "last_mm_tile < 0 (at update)") + + last_m += m + # last_m can be zero if group 0 is skipped + tl.device_assert(last_m >= 0, "last_m < 0 (at update)") + tl.device_assert(last_m <= M, "last_m > M (at update)") + + +# Persistent TGMM kernel. +# ------------------------------------------------------------------------------ + + +@triton.jit +def tgmm_persistent_kernel( + # Tensor pointers: + lhs_ptr, + rhs_ptr, + group_sizes_ptr, + out_ptr, + bias_grad_ptr, + # Tensor shapes: + M: int, + K: int, + N: int, + G: int, + # Meta-parameters: + TRANS_LHS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + GROUP_SIZE: tl.constexpr, + GRID_DIM: tl.constexpr, + COMPUTE_BIAS_GRAD: tl.constexpr, + ACCUMULATE: tl.constexpr, +): + tl.assume(M > 0) + tl.assume(K > 0) + tl.assume(N > 0) + tl.assume(G > 0) + + num_k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + tl.device_assert(num_k_tiles > 0, "num_k_tiles <= 0") + + num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) + tl.device_assert(num_n_tiles > 0, "num_n_tiles <= 0") + + num_tiles = num_k_tiles * num_n_tiles + tl.device_assert(num_tiles > 0, "num_tiles <= 0") + + # Current tile. Each program computes multiple tiles of each group. + tile = tl.program_id(0) + tl.device_assert(tile >= 0, "tile < 0 (at initialization)") + + # Tile limit of last MM problem (inclusive). + last_mm_tile = 0 + + # Last input column of lhs and input row of rhs. Each group reads some + # columns of lhs and some rows of rhs. + last_m = 0 + + # Loop through all (K, m, N) MM problems: + # (K, m) x (m, N) = (K, N) + # sum(m) = M + for g in range(G): + # Get m dimension of current MM problem. + m = tl.load(group_sizes_ptr + g) + # m can be zero if group is empty + tl.device_assert(m >= 0, "m < 0") + + # Loop through tiles of current MM problem. + while tile >= last_mm_tile and tile < last_mm_tile + num_tiles: + # Figure out tile coordinates in current MM problem. + tile_in_mm = tile - last_mm_tile + tl.device_assert(tile_in_mm >= 0, "tile_in_mm < 0") + + tile_k, tile_n = _remap_xcd_tile_grid( + tile_in_mm, num_k_tiles, num_n_tiles, GROUP_SIZE=GROUP_SIZE + ) + + # Do regular MM: + + tl.device_assert(tile_k * BLOCK_SIZE_K >= 0, "tile_k * BLOCK_SIZE_K < 0") + tl.device_assert(tile_n * BLOCK_SIZE_N >= 0, "tile_n * BLOCK_SIZE_N < 0") + + offs_lhs_k = ( + tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + ) % K + offs_rhs_n = ( + tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + ) % N + offs_m = tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + + if TRANS_LHS: + lhs_ptrs = ( + lhs_ptr + offs_lhs_k[:, None] + (last_m + offs_m[None, :]) * K + ) + else: + lhs_ptrs = ( + lhs_ptr + offs_lhs_k[:, None] * M + (last_m + offs_m[None, :]) + ) + + rhs_ptrs = rhs_ptr + (last_m + offs_m[:, None]) * N + offs_rhs_n[None, :] + + loop_m = tl.cdiv(m, BLOCK_SIZE_M) + m_divisible_by_block_m = m % BLOCK_SIZE_M == 0 + if not m_divisible_by_block_m: + loop_m -= 1 + + acc = tl.zeros((BLOCK_SIZE_K, BLOCK_SIZE_N), dtype=tl.float32) + + # Initialize bias accumulator + bias_acc = tl.zeros((BLOCK_SIZE_K,), dtype=tl.float32) + + for _ in range(0, loop_m): + lhs = tl.load(lhs_ptrs) + rhs = tl.load(rhs_ptrs) + + acc = tl.dot(lhs, rhs, acc=acc) + + # Accumulate for bias gradient: sum lhs across M dimension + if COMPUTE_BIAS_GRAD and tile_n == 0: + bias_acc += tl.sum( + lhs, axis=1 + ) # Sum across M dimension [K, M] -> [K] + + if TRANS_LHS: + lhs_ptrs += BLOCK_SIZE_M * K + else: + lhs_ptrs += BLOCK_SIZE_M + + rhs_ptrs += BLOCK_SIZE_M * N + + if not m_divisible_by_block_m: + offs_lhs_k = ( + tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + ) % K + offs_rhs_n = ( + tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + ) % N + offs_m = loop_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + lhs = tl.load(lhs_ptrs, mask=offs_m[None, :] < m, other=0) + rhs = tl.load(rhs_ptrs, mask=offs_m[:, None] < m, other=0) + acc = tl.dot(lhs, rhs, acc=acc) + + # Accumulate last chunk for bias gradient + if COMPUTE_BIAS_GRAD and tile_n == 0: + bias_acc += tl.sum(lhs, axis=1) + + acc = acc.to(out_ptr.type.element_ty) + + offs_out_k = tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + offs_out_n = tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + + out_ptrs = ( + out_ptr + + g.to(tl.int64) * K * N + + offs_out_k[:, None] * N + + offs_out_n[None, :] + ) + + mask = (offs_out_k[:, None] < K) & (offs_out_n[None, :] < N) + if ACCUMULATE: + # Load existing values and add to them (like beta=1 in BLAS) + old_vals = tl.load(out_ptrs, mask=mask, other=0.0) + tl.store(out_ptrs, acc + old_vals, mask=mask) + else: + # Overwrite output (like beta=0 in BLAS) + tl.store(out_ptrs, acc, mask=mask) + + # Store bias gradient (only for first N tile, sum across all M) + if COMPUTE_BIAS_GRAD and tile_n == 0: + # Keep as float32 for atomic_add (bf16 not supported for atomics) + bias_grad_ptrs = bias_grad_ptr + g.to(tl.int64) * K + offs_out_k + # Use atomic add since multiple K-tiles may write to same expert's bias + tl.atomic_add( + bias_grad_ptrs, bias_acc, mask=offs_out_k < K, sem="relaxed" + ) + + # Go to the next tile by advancing number of programs. + tile += GRID_DIM + tl.device_assert(tile > 0, "tile <= 0 (at update)") + + # Get ready to go to the next MM problem. + + last_mm_tile += num_tiles + # last_mm_tile can be zero if group 0 is skipped + tl.device_assert(last_mm_tile >= 0, "last_mm_tile < 0 (at update)") + + last_m += m + # last_m can be zero if group 0 is skipped + tl.device_assert(last_m >= 0, "last_m < 0 (at update)") + tl.device_assert(last_m <= M, "last_m > M (at update)") + + +# Regular non-persistent TGMM kernel. +# ------------------------------------------------------------------------------ + + +@triton.heuristics({"BLOCK_SIZE_G": lambda META: triton.next_power_of_2(META["G"])}) +@triton.jit +def tgmm_non_persistent_kernel( + # Tensor pointers: + lhs_ptr, + rhs_ptr, + group_sizes_ptr, + out_ptr, + bias_grad_ptr, + # Tensor shapes: + M: int, + K: int, + N: int, + G: int, + # Meta-parameters: + TRANS_LHS: tl.constexpr, + BLOCK_SIZE_G: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + GROUP_SIZE: tl.constexpr, + COMPUTE_BIAS_GRAD: tl.constexpr, + ACCUMULATE: tl.constexpr, +): + tl.assume(M > 0) + tl.assume(K > 0) + tl.assume(N > 0) + tl.assume(G > 0) + + # Get group ID from grid. + g = tl.program_id(0) + tl.device_assert(g >= 0, "g < 0") + tl.device_assert(g < G, "g >= G") + + # Get m dimension of current MM group. + m = tl.load(group_sizes_ptr + g) + # m can be zero if group is empty. + tl.device_assert(m >= 0, "m < 0") + + # Skip empty groups. + if m == 0: + return + + # Compute sum(group_sizes) until current group g. + # It's the starting column of lhs and starting row of rhs. + offs_g = tl.arange(0, BLOCK_SIZE_G) + group_sizes = tl.load(group_sizes_ptr + offs_g, mask=offs_g < g, other=0) + start_m = tl.sum(group_sizes) + + num_k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + tl.device_assert(num_k_tiles > 0, "num_k_tiles <= 0") + + num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) + tl.device_assert(num_n_tiles > 0, "num_n_tiles <= 0") + + # Get MM tile from grid. + tile_in_mm = tl.program_id(1) + tl.device_assert(tile_in_mm >= 0, "tile_in_mm < 0") + + tile_k, tile_n = _remap_xcd_tile_grid( + tile_in_mm, num_k_tiles, num_n_tiles, GROUP_SIZE=GROUP_SIZE + ) + + tl.device_assert(tile_k * BLOCK_SIZE_K >= 0, "tile_k * BLOCK_SIZE_K < 0") + tl.device_assert(tile_n * BLOCK_SIZE_N >= 0, "tile_n * BLOCK_SIZE_N < 0") + + offs_lhs_k = (tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)) % K + offs_rhs_n = (tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_m = tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + + if TRANS_LHS: + lhs_ptrs = lhs_ptr + offs_lhs_k[:, None] + (start_m + offs_m[None, :]) * K + else: + lhs_ptrs = lhs_ptr + offs_lhs_k[:, None] * M + (start_m + offs_m[None, :]) + + rhs_ptrs = rhs_ptr + (start_m + offs_m[:, None]) * N + offs_rhs_n[None, :] + + loop_m = tl.cdiv(m, BLOCK_SIZE_M) + m_divisible_by_block_m = m % BLOCK_SIZE_M == 0 + if not m_divisible_by_block_m: + loop_m -= 1 + + acc = tl.zeros((BLOCK_SIZE_K, BLOCK_SIZE_N), dtype=tl.float32) + # Initialize bias accumulator + bias_acc = tl.zeros((BLOCK_SIZE_K,), dtype=tl.float32) + + for _ in range(0, loop_m): + lhs = tl.load(lhs_ptrs) + rhs = tl.load(rhs_ptrs) + + acc = tl.dot(lhs, rhs, acc=acc) + + # Accumulate for bias gradient: sum lhs across M dimension + if COMPUTE_BIAS_GRAD and tile_n == 0: + bias_acc += tl.sum(lhs, axis=1) # [K, M] -> [K] + + if TRANS_LHS: + lhs_ptrs += BLOCK_SIZE_M * K + else: + lhs_ptrs += BLOCK_SIZE_M + + rhs_ptrs += BLOCK_SIZE_M * N + + if not m_divisible_by_block_m: + offs_lhs_k = ( + tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + ) % K + offs_rhs_n = ( + tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + ) % N + offs_m = loop_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + lhs = tl.load(lhs_ptrs, mask=offs_m[None, :] < m, other=0) + rhs = tl.load(rhs_ptrs, mask=offs_m[:, None] < m, other=0) + acc = tl.dot(lhs, rhs, acc=acc) + # Accumulate last chunk for bias gradient + if COMPUTE_BIAS_GRAD and tile_n == 0: + bias_acc += tl.sum(lhs, axis=1) + + acc = acc.to(out_ptr.type.element_ty) + + offs_out_k = tile_k.to(tl.int64) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + offs_out_n = tile_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + + out_ptrs = ( + out_ptr + g.to(tl.int64) * K * N + offs_out_k[:, None] * N + offs_out_n[None, :] + ) + + mask = (offs_out_k[:, None] < K) & (offs_out_n[None, :] < N) + if ACCUMULATE: + # Load existing values and add to them (like beta=1 in BLAS) + old_vals = tl.load(out_ptrs, mask=mask, other=0.0) + tl.store(out_ptrs, acc + old_vals, mask=mask) + else: + # Overwrite output (like beta=0 in BLAS) + tl.store(out_ptrs, acc, mask=mask) + + # Store bias gradient (only for first N tile, sum across all M) + if COMPUTE_BIAS_GRAD and tile_n == 0: + # Keep as float32 for atomic_add (bf16/fp16 not supported for atomics) + bias_grad_ptrs = bias_grad_ptr + g.to(tl.int64) * K + offs_out_k + # Use atomic add since multiple K-tiles may write to same expert's bias + tl.atomic_add(bias_grad_ptrs, bias_acc, mask=offs_out_k < K, sem="relaxed") diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/adapter.py b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..98c224244f27445384e0c2377d73516406927536 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/adapter.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Adapt AITER's Triton grouped GEMM to MegaBlocks' ``gmm`` calling convention. + +MegaBlocks (following tgale96/grouped_gemm) uses a single ``gmm`` entry point +with ``trans_a`` / ``trans_b`` flags: + +* ``trans_a=False, trans_b=False``: a(M,K) @ b(G,K,N) -> c(M,N) +* ``trans_a=False, trans_b=True`` : a(M,K) @ b(G,N,K)^T -> c(M,N) (dgrad) +* ``trans_a=True`` : a(M,K)^T @ b(M,N) per group -> c(G,K,N) (wgrad) + +AITER exposes these as two kernels: ``gmm`` ((M,K)@(G,K,N)->(M,N), transposition +of the 3D operand inferred from strides) and ``ptgmm`` ((K,M)@(M,N)->(G,K,N), +transposition of the 2D operand inferred from strides). +""" + +import torch + +from .gmm import gmm as _aiter_gmm +from .gmm import ptgmm as _aiter_ptgmm + + +def gmm(a, b, c, batch_sizes, trans_a=False, trans_b=False): + # AITER requires group sizes to be int32 and to live on the compute device. + group_sizes = batch_sizes.to(device=a.device, dtype=torch.int32) + + # AITER asserts exact strides: gmm wants lhs/rhs row-major (a transposed + # 3D operand must be exactly column-major), tgmm wants rhs row-major and + # lhs row/column-major. Make operands contiguous first so the transposed + # views have the precise strides the kernels expect. `.contiguous()` is a + # no-op when the tensor is already contiguous. + if trans_a: + # Weight gradient: a(M,K), b(M,N) -> c(G,K,N). + # Pass a transposed so AITER sees lhs(K,M) column-major (TRANS_LHS). + _aiter_ptgmm( + a.contiguous().transpose(0, 1), + b.contiguous(), + group_sizes, + preferred_element_type=c.dtype, + existing_out=c, + ) + else: + # trans_b contracts b's last dim: pass a column-major (G,K,N) view. + rhs = b.contiguous() + if trans_b: + rhs = rhs.transpose(1, 2) + _aiter_gmm( + a.contiguous(), + rhs, + group_sizes, + preferred_element_type=c.dtype, + existing_out=c, + ) + return c diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/configs.py b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/configs.py new file mode 100644 index 0000000000000000000000000000000000000000..9a4fe5617d8100869aa76dba9b7d22c7bcab814f --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/configs.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: MIT +# Tuned GMM configs vendored from ROCm/aiter (aiter/ops/triton/configs/). +# Inlined as a Python module so packaging always includes them. + +CONFIGS = {'gfx1250': {'gmm': {'default': {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_K': 64, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}}, 'ptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}}, 'nptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}}}, 'gfx942': {'gmm': {'default': {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_K': 64, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 304, 'num_warps': 8, 'num_stages': 1}}, 'ptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 304, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'GRID_DIM': 304, 'num_warps': 8, 'num_stages': 1}}, 'nptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}}}, 'gfx950': {'gmm': {'default': {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_K': 64, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}}, 'ptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'GRID_DIM': 256, 'num_warps': 8, 'num_stages': 1}}, 'nptgmm': {'default': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 256, 'BLOCK_SIZE_N': 256, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}, 'accumulate': {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_K': 128, 'BLOCK_SIZE_N': 128, 'GROUP_SIZE': 1, 'num_warps': 8, 'num_stages': 1}}}} diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/gmm.py b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/gmm.py new file mode 100644 index 0000000000000000000000000000000000000000..e30c9326c6d4e4836d1303e2761ea2440a7f4750 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/gmm.py @@ -0,0 +1,567 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved. + + +# Imports. +# ------------------------------------------------------------------------------ + +# PyTorch +import torch +from torch import Tensor + +# Triton +import triton + +# AITER: GMM utility functions +from .utils.gmm_common import ( + DTYPE, + is_power_of_2, + check_input_device_dtype, + check_bias_shape_stride, + get_gmm_shape, + get_gmm_output, + get_gmm_transposition, + get_tgmm_shape, + get_tgmm_output, + get_tgmm_bias_grad, + get_tgmm_transposition, +) + +# AITER: GMM Triton kernels +from ._triton_kernels.gmm import ( + gmm_kernel, + tgmm_persistent_kernel, + tgmm_non_persistent_kernel, + get_config, +) + +# GMM PyTorch wrapper. +# ------------------------------------------------------------------------------ + + +def _gmm_grid( + N: int, + block_size_m: int, + block_size_n: int, + group_sizes: Tensor, + grid_dim: int, +) -> tuple[int]: + assert N > 0, f"N must be positive, it's {N}." + assert is_power_of_2( + block_size_m + ), f"M-dimension tile size must be a power of 2 (it's {block_size_m})." + assert is_power_of_2( + block_size_n + ), f"N-dimension tile size must be a power of 2 (it's {block_size_n})." + assert torch.all(group_sizes >= 0).item(), "All group_sizes must be non-negative." + assert grid_dim > 0, f"Grid dimension must be positive (it's {grid_dim})." + num_m_tiles = (group_sizes + block_size_m - 1) // block_size_m + assert torch.all(num_m_tiles >= 0).item(), "All num_m_tiles must be non-negative." + num_n_tiles = triton.cdiv(N, block_size_n) + assert num_n_tiles > 0, f"num_n_tiles must be positive, it's {num_n_tiles}." + num_tiles = torch.sum(num_m_tiles * num_n_tiles).item() + assert num_tiles > 0, f"num_tiles must be positive, it's {num_tiles}." + num_programs = int(min(grid_dim, num_tiles)) + assert num_programs > 0, f"num_programs must be positive, it's {num_programs}." + return (num_programs,) + + +def gmm( + lhs: Tensor, + rhs: Tensor, + group_sizes: Tensor, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, + config: dict[str, int] | None = None, + bias: Tensor | None = None, +) -> Tensor: + """ + Perform Group Matrix Multiplication (GMM): out = lhs @ rhs + bias + + lhs rows are divided into G groups. Each group of lhs rows is matrix multiplied with a plane of + rhs 3D tensor and then stored in a slice of out. In PyTorch parlance, it can be implemented as + follows for a given group g: + out[group_start:group_end, :] = lhs[group_start:group_end, :] @ rhs[g] + bias[g] + + The size of each group, and their respective start and end positions are specified by + group_sizes tensor. For instance, suppose that group_sizes = [3, 2, 4, 1]. In this particular + case we have 4 groups. The 1st group starts at 0 and ends at 2, the second group starts at 3 and + ends at 4, the third group starts at 5 and ends at 8, and the fourth and final group consists of + just the 10th (last) row of lhs. + + Parameters + ---------- + lhs : torch.Tensor + Left-hand side 2D input tensor. Shape: (M, K). + lhs data type must be torch.float16 or torch.bfloat16, and must match rhs data type. + lhs must be on the same device of rhs and group_sizes. + rhs : torch.Tensor + Right-hand side 3D input tensor. Shape: (G, K, N). + rhs data type must be torch.float16 or torch.bfloat16, and must match lhs data type. + rhs must be on the same device of lhs and group_sizes. + group_sizes : torch.Tensor + 1D input tensor describing group sizes. Shape: (G,). + group_sizes data type must be torch.int32 and all its elements must be non-negative. + group_sizes must be on the same device of lhs and rhs. + preferred_element_type : torch.dtype, optional + Desired data type for output tensor. Default is torch.bfloat16. + Supported output types are torch.float16 and torch.bfloat16. + existing_out : torch.Tensor or None, optional + Preallocated output tensor. Default is None. + If provided, results are written into this tensor. Otherwise, a new output tensor is + allocated. + If provided then it must have shape (M, N), its data type must match preferred_element_type + and it must be on the same device of other input tensors. + config : dict[str, int] or None, optional + Optional dictionary with kernel metaparameters. If absent, config will be queried from + internal tuning database. + bias : torch.Tensor or None, optional + Optional bias tensor. Shape: (G, N). + If provided, bias data type must match lhs and rhs data type, and bias must be on the same + device as other input tensors. Each group g adds bias[g] to the output. + + Returns + ------- + torch.Tensor + The computed output 2D tensor. Shape: (M, N). + Output tensor data type is given by preferred_element_type. + If existing_out is provided then existing_out is also returned. + + Implementation Notes + -------------------- + - GMM is implemented with a persistent Triton kernel. + - lhs must be row-major (lhs.stride() == (K, 1)). + - rhs can be row-major (rhs.stride() == (K * N, N, 1)) or column-major (rhs.stride() == + (K * N, 1, K)). If rhs is row-major then kernel parameter TRANS_RHS == False, this is useful + for implementing forward pass. If rhs is column-major then kernel parameter TRANS_RHS == True, + this is useful for computing the lhs derivative in the backward pass, while fusing the + transposition. + - out must be row-major (out.stride() == (N, 1)). + - bias must be row-major (bias.stride() == (N, 1)) if provided. + """ + use_bias = bias is not None + check_input_device_dtype(lhs, rhs, group_sizes, bias) + + M, K, N, G = get_gmm_shape(lhs, rhs, group_sizes) + + if use_bias: + check_bias_shape_stride(bias, G, N) + + out = get_gmm_output( + M, + N, + device=lhs.device, + preferred_element_type=preferred_element_type, + existing_out=existing_out, + ) + + trans_rhs, _ = get_gmm_transposition(lhs, rhs, out) + + if config is None: + config = get_config("gmm", M, K, N, G) + + assert all( + key in config + and isinstance(config[key], int) + and ( + is_power_of_2(config[key]) + if key.startswith("BLOCK_SIZE_") + else config[key] > 0 + ) + for key in { + "BLOCK_SIZE_M", + "BLOCK_SIZE_K", + "BLOCK_SIZE_N", + "GROUP_SIZE", + "GRID_DIM", + } + ), "Invalid GMM kernel config." + + grid = _gmm_grid( + N, + config["BLOCK_SIZE_M"], + config["BLOCK_SIZE_N"], + group_sizes, + config["GRID_DIM"], + ) + + # fmt: off + gmm_kernel[grid]( + # Tensor pointers: + lhs, rhs, group_sizes, out, bias, + # Tensor shapes: + M, K, N, G, + # Meta-parameters: + TRANS_RHS=trans_rhs, + USE_BIAS=use_bias, + **config, + ) + # fmt: on + + return out + + +# Persistent TGMM PyTorch wrapper. +# ------------------------------------------------------------------------------ + + +def _ptgmm_grid( + K: int, + N: int, + G: int, + block_size_k: int, + block_size_n: int, + grid_dim: int, +) -> tuple[int]: + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}." + assert G > 0, f"G must be positive, it's {G}." + assert is_power_of_2( + block_size_k + ), f"K-dimension tile size must be a power of 2 (it's {block_size_k})." + assert is_power_of_2( + block_size_n + ), f"N-dimension tile size must be a power of 2 (it's {block_size_n})." + assert grid_dim > 0, f"Grid dimension must be positive (it's {grid_dim})." + num_k_tiles = triton.cdiv(K, block_size_k) + assert num_k_tiles > 0, f"num_k_tiles must be positive, it's {num_k_tiles}." + num_n_tiles = triton.cdiv(N, block_size_n) + assert num_n_tiles > 0, f"num_n_tiles must be positive, it's {num_n_tiles}." + num_tiles = G * num_k_tiles * num_n_tiles + assert num_tiles > 0, f"num_tiles must be positive, it's {num_tiles}." + num_programs = min(grid_dim, num_tiles) + assert num_programs > 0, f"num_programs must be positive, it's {num_programs}." + return (num_programs,) + + +def ptgmm( + lhs: Tensor, + rhs: Tensor, + group_sizes: Tensor, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, + config: dict[str, int] | None = None, + bias_grad: Tensor | None = None, + accumulate: bool = False, +) -> Tensor: + """ + Perform a Group Matrix Multiplication (GMM) variant: out = lhs @ rhs + + lhs columns and rhs rows are divided into G groups. Each group of lhs is matrix multiplied with + the respective group of rhs and then stored in a plane of the output 3D tensor. In PyTorch + parlance, it can be implemented as follows for a given group g: + out[g] = lhs[:, group_start:group_end] @ rhs[group_start:group_end, :] + + The 't' in the operator name derives from MaxText implementation + (https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/kernels/megablox/gmm.py), + which served as the initial inspiration for this one. TGMM differs from GMM in terms of tensor + shapes. GMM does (M, K) @ (G, K, N) = (M, N) while TGMM does (K, M) @ (M, N) = (G, K, N). + + The 'p' in the operator name means that it is implemented with a persistent kernel. There is + also the non-persistent variation, which is implemented with a regular kernel. Please take a + look at nptgmm operator. Both ptgmm and nptgmm implement the same computation, choosing one or + the other is a matter of performance for the target workload. + + Parameters + ---------- + lhs : torch.Tensor + Left-hand side 2D input tensor. Shape: (K, M). + lhs data type must be torch.float16 or torch.bfloat16, and must match rhs data type. + lhs must be on the same device of rhs and group_sizes. + rhs : torch.Tensor + Right-hand side 2D input tensor. Shape: (M, N). + rhs data type must be torch.float16 or torch.bfloat16, and must match lhs data type. + rhs must be on the same device of lhs and group_sizes. + group_sizes : torch.Tensor + 1D input tensor describing group sizes. Shape: (G,). + group_sizes data type must be torch.int32 and all its elements must be non-negative. + group_sizes must be on the same device of lhs and rhs. + preferred_element_type : torch.dtype, optional + Desired data type for output tensor. Default is torch.bfloat16. + Supported output types are torch.float16 and torch.bfloat16. + existing_out : torch.Tensor or None, optional + Preallocated output tensor. Default is None. + If provided, results are written into this tensor. Otherwise, a new output tensor is + allocated. + If provided then it must have shape (G, K, N), its data type must match + preferred_element_type and it must be on the same device of other input tensors. + config : dict[str, int] or None, optional + Optional dictionary with kernel metaparameters. If absent, config will be queried from + internal tuning database. + bias_grad : torch.Tensor or None, optional + Optional bias gradient output tensor. Shape: (G, K). + If provided, the kernel will compute the bias gradient and write it to this tensor. + bias_grad must be torch.float32 (kernel uses atomic_add which requires float32), + accumulate : bool, optional + Whether to accumulate into existing output tensor values. Default is False. + If False, output will be overwritten with fresh computation. + If True, results will be added to existing output tensor values. + + Returns + ------- + torch.Tensor + The computed output 3D tensor. Shape: (G, K, N). + Output tensor data type is given by preferred_element_type. + If existing_out is provided then existing_out is also returned. + + Implementation Notes + -------------------- + - PTGMM is implemented with a persistent Triton kernel. + - lhs can be row-major (lhs.stride() == (M, 1)) or column-major (lhs.stride() == (1, K)). If lhs + is row-major then kernel parameter TRANS_LHS == False. If lhs is column-major then kernel + parameter TRANS_LHS == True, this is useful for computing the rhs derivative in the backward + pass, while fusing the transposition. + - rhs must be row-major (rhs.stride() == (N, 1)). + - out must be row-major (out.stride() == (K * N, N, 1)). + """ + check_input_device_dtype(lhs, rhs, group_sizes) + + M, K, N, G = get_tgmm_shape(lhs, rhs, group_sizes) + + out = get_tgmm_output( + K, + N, + G, + device=lhs.device, + preferred_element_type=preferred_element_type, + existing_out=existing_out, + ) + + trans_lhs, _ = get_tgmm_transposition(lhs, rhs, out) + + if config is None: + config = get_config("ptgmm", M, K, N, G, accumulate) + + assert all( + key in config + and isinstance(config[key], int) + and ( + is_power_of_2(config[key]) + if key.startswith("BLOCK_SIZE_") + else config[key] > 0 + ) + for key in { + "BLOCK_SIZE_M", + "BLOCK_SIZE_K", + "BLOCK_SIZE_N", + "GROUP_SIZE", + "GRID_DIM", + } + ), "Invalid PTGMM kernel config." + + # Bias gradient handling. + # ----------------------- + # Get or validate bias gradient tensor. + compute_bias_grad = bias_grad is not None + bias_grad_ptr = get_tgmm_bias_grad( + K, + G, + device=lhs.device, + existing_bias_grad=bias_grad, + ) + + grid = _ptgmm_grid( + K, + N, + G, + config["BLOCK_SIZE_K"], + config["BLOCK_SIZE_N"], + config["GRID_DIM"], + ) + + # fmt: off + tgmm_persistent_kernel[grid]( + # Tensor pointers: + lhs, rhs, group_sizes, out, bias_grad_ptr, + # Tensor shapes: + M, K, N, G, + # Meta-parameters: + TRANS_LHS=trans_lhs, + COMPUTE_BIAS_GRAD=compute_bias_grad, + ACCUMULATE=accumulate, + **config, + ) + # fmt: on + + return out + + +# Regular non-persistent TGMM PyTorch wrapper. +# ------------------------------------------------------------------------------ + + +def _nptgmm_grid( + K: int, + N: int, + G: int, + block_size_k: int, + block_size_n: int, +) -> tuple[int, int]: + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}." + assert G > 0, f"G must be positive, it's {G}." + assert is_power_of_2( + block_size_k + ), f"K-dimension tile size must be a power of 2 (it's {block_size_k})." + assert is_power_of_2( + block_size_n + ), f"N-dimension tile size must be a power of 2 (it's {block_size_n})." + num_k_tiles = triton.cdiv(K, block_size_k) + assert num_k_tiles > 0, f"num_k_tiles must be positive, it's {num_k_tiles}." + num_n_tiles = triton.cdiv(N, block_size_n) + assert num_n_tiles > 0, f"num_n_tiles must be positive, it's {num_n_tiles}." + num_tiles_per_mm = num_k_tiles * num_n_tiles + assert ( + num_tiles_per_mm > 0 + ), f"num_tiles_per_mm must be positive, it's {num_tiles_per_mm}." + return (G, num_tiles_per_mm) + + +def nptgmm( + lhs: Tensor, + rhs: Tensor, + group_sizes: Tensor, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, + config: dict[str, int] | None = None, + bias_grad: Tensor | None = None, + accumulate: bool = False, +) -> Tensor: + """ + Perform a Group Matrix Multiplication (GMM) variant: out = lhs @ rhs + + lhs columns and rhs rows are divided into G groups. Each group of lhs is matrix multiplied with + the respective group of rhs and then stored in a plane of the output 3D tensor. In PyTorch + parlance, it can be implemented as follows for a given group g: + out[g] = lhs[:, group_start:group_end] @ rhs[group_start:group_end, :] + + The 't' in the operator name derives from MaxText implementation + (https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/kernels/megablox/gmm.py), + which served as the initial inspiration for this one. TGMM differs from GMM in terms of tensor + shapes. GMM does (M, K) @ (G, K, N) = (M, N) while TGMM does (K, M) @ (M, N) = (G, K, N). + + The 'np' in the operator name means that it is implemented with a non-persistent, i.e. regular + kernel. There is also the persistent variation, which is implemented with a persistent kernel. + Please take a look at ptgmm operator. Both nptgmm and ptgmm implement the same computation, + choosing one or the other is a matter of performance for the target workload. + + Parameters + ---------- + lhs : torch.Tensor + Left-hand side 2D input tensor. Shape: (K, M). + lhs data type must be torch.float16 or torch.bfloat16, and must match rhs data type. + lhs must be on the same device of rhs and group_sizes. + rhs : torch.Tensor + Right-hand side 2D input tensor. Shape: (M, N). + rhs data type must be torch.float16 or torch.bfloat16, and must match lhs data type. + rhs must be on the same device of lhs and group_sizes. + group_sizes : torch.Tensor + 1D input tensor describing group sizes. Shape: (G,). + group_sizes data type must be torch.int32 and all its elements must be non-negative. + group_sizes must be on the same device of lhs and rhs. + preferred_element_type : torch.dtype, optional + Desired data type for output tensor. Default is torch.bfloat16. + Supported output types are torch.float16 and torch.bfloat16. + existing_out : torch.Tensor or None, optional + Preallocated output tensor. Default is None. + If provided, results are written into this tensor. Otherwise, a new output tensor is + allocated. + If provided then it must have shape (G, K, N), its data type must match + preferred_element_type and it must be on the same device of other input tensors. + config : dict[str, int] or None, optional + Optional dictionary with kernel metaparameters. If absent, config will be queried from + internal tuning database. + bias_grad : torch.Tensor or None, optional + Optional bias gradient output tensor. Shape: (G, K). + If provided, the kernel will compute the bias gradient and write it to this tensor. + bias_grad must be torch.float32 (kernel uses atomic_add which requires float32), + accumulate : bool, optional + Whether to accumulate into existing output tensor values. Default is False. + If False, output will be overwritten with fresh computation. + If True, results will be added to existing output tensor values. + + Returns + ------- + torch.Tensor + The computed output 3D tensor. Shape: (G, K, N). + Output tensor data type is given by preferred_element_type. + If existing_out is provided then existing_out is also returned. + + Implementation Notes + -------------------- + - NPTGMM is implemented with a non-persistent regular Triton kernel. + - lhs can be row-major (lhs.stride() == (M, 1)) or column-major (lhs.stride() == (1, K)). If lhs + is row-major then kernel parameter TRANS_LHS == False. If lhs is column-major then kernel + parameter TRANS_LHS == True, this is useful for computing the rhs derivative in the backward + pass, while fusing the transposition. + - rhs must be row-major (rhs.stride() == (N, 1)). + - out must be row-major (out.stride() == (K * N, N, 1)). + """ + check_input_device_dtype(lhs, rhs, group_sizes) + + M, K, N, G = get_tgmm_shape(lhs, rhs, group_sizes) + + out = get_tgmm_output( + K, + N, + G, + device=lhs.device, + preferred_element_type=preferred_element_type, + existing_out=existing_out, + ) + + trans_lhs, _ = get_tgmm_transposition(lhs, rhs, out) + + # Bias gradient handling. + # ----------------------- + # Get or validate bias gradient tensor. + compute_bias_grad = bias_grad is not None + bias_grad_ptr = get_tgmm_bias_grad( + K, + G, + device=lhs.device, + existing_bias_grad=bias_grad, + ) + + if config is None: + config = get_config("nptgmm", M, K, N, G, accumulate) + + assert all( + key in config + and isinstance(config[key], int) + and ( + is_power_of_2(config[key]) + if key.startswith("BLOCK_SIZE_") + else config[key] > 0 + ) + for key in { + "BLOCK_SIZE_M", + "BLOCK_SIZE_K", + "BLOCK_SIZE_N", + "GROUP_SIZE", + } + ), "Invalid NPTGMM kernel config." + + grid = _nptgmm_grid( + K, + N, + G, + config["BLOCK_SIZE_K"], + config["BLOCK_SIZE_N"], + ) + + # fmt: off + tgmm_non_persistent_kernel[grid]( + # Tensor pointers: + lhs, rhs, group_sizes, out, bias_grad_ptr, + # Tensor shapes: + M, K, N, G, + # Meta-parameters: + TRANS_LHS=trans_lhs, + COMPUTE_BIAS_GRAD=compute_bias_grad, + ACCUMULATE=accumulate, + **config, + ) + # fmt: on + + return out diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/_triton/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/_triton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/_triton/arch_info.py b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/_triton/arch_info.py new file mode 100644 index 0000000000000000000000000000000000000000..3f6c88581a64044518125623f116082c53bd5474 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/_triton/arch_info.py @@ -0,0 +1,46 @@ +import triton + +# Detect the GPU arch lazily: querying the triton driver at import time fails +# in headless environments (e.g. the kernel-builder ABI check sandbox has no +# GPU), and the original JAX fallback pulled in an unrelated runtime dep. The +# arch is only actually needed when a GMM kernel is dispatched, so resolve and +# cache on first call. +_CACHED_ARCH = None + + +def get_arch(): + global _CACHED_ARCH + if _CACHED_ARCH is not None: + return _CACHED_ARCH + try: + _CACHED_ARCH = triton.runtime.driver.active.get_current_target().arch + except RuntimeError: + try: + from jax._src.lib import gpu_triton as triton_kernel_call_lib + _CACHED_ARCH = triton_kernel_call_lib.get_arch_details("0").split(":")[0] + except ImportError as e: + raise RuntimeError( + "Cannot determine GPU arch: triton driver is inactive and " + "JAX is not available. A GPU is required for grouped GEMM." + ) from e + return _CACHED_ARCH + + +def is_gluon_avail(): + return get_arch() in ("gfx950", "gfx1250") + + +def is_fp4_avail(): + return get_arch() in ("gfx950", "gfx1250") + + +def is_fp8_avail(): + return get_arch() in ("gfx942", "gfx950", "gfx1250", "gfx1200", "gfx1201") + + +def is_mx_scale_preshuffling_avail(): + return get_arch() in ("gfx950", "gfx1250") + + +def is_tdm_avail(): + return get_arch() in ("gfx1250",) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/_triton/pid_preprocessing.py b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/_triton/pid_preprocessing.py new file mode 100644 index 0000000000000000000000000000000000000000..99792bb3ba2fab8fff223bba733ced1eb6e6df53 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/_triton/pid_preprocessing.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: MIT + +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import triton +import triton.language as tl + + +@triton.jit +def remap_xcd_chunked( + pid, GRID_MN, NUM_XCDS: tl.constexpr = 8, CHUNK_SIZE: tl.constexpr = 2 +): + # Compute current XCD and local PID + xcd = pid % NUM_XCDS + # distribute the modulo pids in round robin + if pid > (GRID_MN // (NUM_XCDS * CHUNK_SIZE)) * (NUM_XCDS * CHUNK_SIZE): + return pid + local_pid = pid // NUM_XCDS + # Calculate chunk index and position within chunk + chunk_idx = local_pid // CHUNK_SIZE + pos_in_chunk = local_pid % CHUNK_SIZE + # Calculate new PID + new_pid = chunk_idx * NUM_XCDS * CHUNK_SIZE + xcd * CHUNK_SIZE + pos_in_chunk + return new_pid + + +@triton.jit +def remap_xcd(pid, GRID_MN, NUM_XCDS: tl.constexpr = 8): + ## pid remapping on xcds + # Number of pids per XCD in the new arrangement + pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS + # When GRID_MN cannot divide NUM_XCDS, some xcds will have + # pids_per_xcd pids, the other will have pids_per_xcd - 1 pids. + # We calculate the number of xcds that have pids_per_xcd pids as + # tall_xcds + tall_xcds = GRID_MN % NUM_XCDS + tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds + # Compute current XCD and local pid within the XCD + xcd = pid % NUM_XCDS + local_pid = pid // NUM_XCDS + # Calculate new pid based on the new grouping + # Note that we need to consider the following two cases: + # 1. the current pid is on a tall xcd + # 2. the current pid is on a short xcd + if xcd < tall_xcds: + pid = xcd * pids_per_xcd + local_pid + else: + pid = ( + tall_xcds * pids_per_xcd + + (xcd - tall_xcds) * (pids_per_xcd - 1) + + local_pid + ) + + return pid + + +@triton.jit +def pid_grid(pid: int, num_pid_m: int, num_pid_n: int, GROUP_SIZE_M: tl.constexpr = 1): + """ + Maps 1D pid to 2D grid coords (pid_m, pid_n). + + Args: + - pid: 1D pid + - num_pid_m: grid m size + - num_pid_n: grid n size + - GROUP_SIZE_M: tl.constexpr: default is 1 + """ + if GROUP_SIZE_M == 1: + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + else: + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + tl.assume(group_size_m >= 0) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + return pid_m, pid_n + + +@triton.jit +def pid_grid_3d(pid: int, num_pid_m: int, num_pid_n: int, num_pid_k): + """ + Maps 1D pid to 3D grid coords (pid_m, pid_n, pid_k). + Args: + - pid: 1D pid + - num_pid_m: grid m size + - num_pid_n: grid n size + - num_pid_k: grid k size + + Returns: + - pid_m, pid_n, pid_k: 3D grid coordinates + """ + pid_m = pid % num_pid_m + pid_n = (pid // num_pid_m) % num_pid_n + pid_k = pid // (num_pid_m * num_pid_n) % num_pid_k + + return pid_m, pid_n, pid_k diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/gmm_common.py b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/gmm_common.py new file mode 100644 index 0000000000000000000000000000000000000000..153dee65b50ab5f833262481889d2184d1ca639f --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/gmm_common.py @@ -0,0 +1,752 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved. + +# Imports. +# ------------------------------------------------------------------------------ + +# PyTorch +import torch +from torch import Tensor + +# AITER: logging +from .logger import AiterTritonLogger + +_LOGGER: AiterTritonLogger = AiterTritonLogger() + + +# Supported data types. +# ------------------------------------------------------------------------------ + +# Supported data types, as strings. +SUPPORTED_DTYPES_STR: set[str] = {"fp16", "bf16"} + + +# Convert string data type to PyTorch data type. +def dtype_from_str(dtype_str: str) -> torch.dtype: + dtype_str = dtype_str.strip().lower() + dtype_str = dtype_str[1:] if dtype_str[0] in {"i", "o"} else dtype_str + assert ( + dtype_str in SUPPORTED_DTYPES_STR + ), "String data type isn't in set of supported string data types." + return {"fp16": torch.float16, "bf16": torch.bfloat16}[dtype_str] + + +# Supported data types, as PyTorch types. +SUPPORTED_DTYPES: set[torch.dtype] = { + dtype_from_str(dtype_str) for dtype_str in SUPPORTED_DTYPES_STR +} + + +# Convert PyTorch data type to string data type. +def str_from_dtype(dtype: torch.dtype) -> str: + assert ( + dtype in SUPPORTED_DTYPES + ), "PyTorch data type isn't in set of supported PyTorch data types." + return {torch.float16: "fp16", torch.bfloat16: "bf16"}[dtype] + + +# Default data type, as string. +DTYPE_STR: str = "bf16" +assert ( + DTYPE_STR in SUPPORTED_DTYPES_STR +), "Default string data type isn't in set of supported string data types." + + +# Default data type, as PyTorch type. +DTYPE: torch.dtype = dtype_from_str(DTYPE_STR) + + +# Other defaults. +# ------------------------------------------------------------------------------ + +# Default device. +DEVICE: torch.device | str = "cuda" + +# Default RNG seed for input generation. +RNG_SEED: int = 0 + +# Default number of group sizes. +NUM_GROUP_SIZES: int = 1 + +# Default transposition (NN). +TRANS_LHS: bool = False +TRANS_RHS: bool = False + + +# Parameter checking functions. +# ------------------------------------------------------------------------------ + + +def is_power_of_2(x: int) -> bool: + return (x > 0) and (x & (x - 1) == 0) + + +def check_input_device_dtype( + lhs: Tensor, rhs: Tensor, group_sizes: Tensor, bias: Tensor | None = None +) -> None: + assert ( + lhs.device == rhs.device == group_sizes.device + ), f"All input tensors must be in the same device (lhs = {lhs.device}, rhs = {rhs.device}, group_sizes = {group_sizes.device})." + assert ( + lhs.dtype == rhs.dtype + ), f"lhs and rhs types must match (lhs = {lhs.dtype}, rhs = {rhs.dtype})." + assert group_sizes.dtype == torch.int32, "group_sizes type must be int32." + + if bias is not None: + assert ( + bias.device == lhs.device + ), f"bias must be on the same device as lhs (bias = {bias.device}, lhs = {lhs.device})." + assert ( + bias.dtype == lhs.dtype + ), f"bias dtype must match lhs dtype (bias = {bias.dtype}, lhs = {lhs.dtype})." + + +def check_bias_shape_stride(bias: Tensor, G: int, N: int) -> None: + assert bias.shape == ( + G, + N, + ), f"bias must have shape (G, N) = ({G}, {N}), got {bias.shape}." + assert bias.stride() == (N, 1), "bias must be row-major (bias.stride() == (N, 1))." + + +# Generation of group sizes. +# ------------------------------------------------------------------------------ + + +# Probabilities for generating random group sizes. +UNUSED_TOKENS_PROB: float = 0.0 +UNUSED_EXPERTS_PROB: float = 0.1 + + +def gen_uniform_group_sizes( + M: int, + G: int, + device: torch.device | str = DEVICE, +) -> Tensor: + assert M >= 0, f"Number of tokens M must be non-negative (it's {M})." + assert G > 0, f"Number of experts G must be positive (it's {G})." + + base = M // G + remainder = M % G + group_sizes = torch.full((G,), base, dtype=torch.int32, device=device) + if remainder > 0: + group_sizes[:remainder] += 1 + + assert ( + len(group_sizes) == G + ), f"Group sizes don't have {G} elements (it's {len(group_sizes)})." + assert torch.all(group_sizes >= 0).item(), "All group sizes must be non-negative." + assert ( + torch.sum(group_sizes).item() == M + ), f"Group sizes don't add up to total tokens {M}." + assert group_sizes.dtype == torch.int32, "Group sizes must be int32." + + return group_sizes + + +def gen_group_sizes( + M: int, + G: int, + device: torch.device | str = DEVICE, + rng_seed: int | None = RNG_SEED, + unused_tokens_prob: float = UNUSED_TOKENS_PROB, + unused_experts_prob: float = UNUSED_EXPERTS_PROB, +) -> Tensor: + assert M >= 0, f"Number of tokens M must be non-negative (it's {M})." + assert G > 0, f"Number of experts G must be positive (it's {G})." + assert ( + 0 <= unused_tokens_prob <= 1 + ), f"Probability of unused tokens must be in [0, 1] interval (it's {unused_tokens_prob})." + assert ( + 0 <= unused_experts_prob <= 1 + ), f"Probability of unused experts must be in [0, 1] interval (it's {unused_experts_prob})." + + if rng_seed is not None: + torch.manual_seed(rng_seed) + + if unused_tokens_prob > 0: + # Optionally drop tokens to simulate routing sparsity, some tokens may not be routed. + num_unused_tokens = M + while num_unused_tokens == M: + num_unused_tokens = int( + torch.binomial( + torch.tensor(float(M), device=device), + torch.tensor(unused_tokens_prob, device=device), + ).item() + ) + else: + num_unused_tokens = 0 + num_used_tokens = M - num_unused_tokens + assert ( + num_unused_tokens >= 0 + ), f"Number of unused tokens must be non-negative (it's {num_unused_tokens})." + assert ( + num_used_tokens > 0 + ), f"Number of used tokens must be positive (it's {num_used_tokens})." + assert ( + num_used_tokens + num_unused_tokens == M + ), f"Unused + used tokens don't add up total tokens ({num_used_tokens} + {num_unused_tokens} != {M})." + + if num_unused_tokens > 0: + _LOGGER.debug( + f"Group sizes generation: dropped {num_unused_tokens} token{'s' if num_unused_tokens > 1 else ''}.", + ) + + if unused_experts_prob > 0: + # Some experts may have zero tokens assigned to them. + num_used_experts = 0 + while num_used_experts == 0: + used_experts = torch.nonzero( + torch.rand((G,), device=device) >= unused_experts_prob + ).squeeze() + num_used_experts = used_experts.numel() + else: + used_experts = torch.arange(0, G, device=device) + num_used_experts = G + num_unused_experts = G - num_used_experts + assert ( + num_unused_experts >= 0 + ), f"Number of unused experts must be non-negative (it's {num_unused_experts})." + assert ( + num_used_experts >= 1 + ), f"At least one expert must be used (it's {num_used_experts})." + assert ( + num_unused_experts + num_used_experts == G + ), f"Unused + used experts don't add up total experts ({num_unused_experts} + {num_used_experts} != {G})." + + if num_unused_experts > 0: + _LOGGER.debug( + f"Group sizes generation: dropped {num_unused_experts} expert{'s' if num_unused_experts > 1 else ''}.", + ) + + group_sizes = torch.bincount( + used_experts[ + torch.randint(low=0, high=num_used_experts, size=(num_used_tokens,)) + ], + minlength=G, + ).to(torch.int32) + + assert ( + len(group_sizes) == G + ), f"Group sizes don't have {G} elements (it's {len(group_sizes)})." + assert torch.all(group_sizes >= 0).item(), "All group sizes must be non-negative." + assert ( + torch.sum(group_sizes).item() == num_used_tokens + ), f"Group sizes don't add up to used tokens {num_used_tokens}." + assert group_sizes.dtype == torch.int32, "Group sizes must be int32." + + return group_sizes + + +def gen_multiple_group_sizes( + num_group_sizes: int, + M: int, + G: int, + device: torch.device | str = DEVICE, + rng_seed: int | None = RNG_SEED, + unused_tokens_prob: float = UNUSED_TOKENS_PROB, + unused_experts_prob: float = UNUSED_EXPERTS_PROB, + group_sizes_0: Tensor | None = None, +) -> list[Tensor]: + assert ( + num_group_sizes > 0 + ), f"Number of group sizes to be generated must be positive, it's {num_group_sizes}." + multiple_group_sizes = [ + gen_group_sizes( + M, + G, + device=device, + rng_seed=rng_seed if g == 0 else None, + unused_tokens_prob=unused_tokens_prob, + unused_experts_prob=unused_experts_prob, + ) + for g in range( + num_group_sizes if group_sizes_0 is None else num_group_sizes - 1 + ) + ] + if group_sizes_0 is not None: + multiple_group_sizes.insert(0, group_sizes_0) + assert ( + len(multiple_group_sizes) == num_group_sizes + ), f"Expecting {num_group_sizes} distinct group sizes (it's {len(multiple_group_sizes)})." + return multiple_group_sizes + + +# GMM helpers: tensor generation. +# ------------------------------------------------------------------------------ + + +def gen_gmm_input( + M: int, + K: int, + N: int, + G: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, + trans_rhs: bool = TRANS_RHS, + rng_seed: int | None = RNG_SEED, + unif_group_sizes: bool = False, +) -> tuple[Tensor, Tensor, Tensor]: + assert M > 0, f"Number of lhs rows M must be positive (M = {M})." + assert K > 0, f"Number of lhs columns / rhs rows K must be positive (K = {K})." + assert N > 0, f"Number of rhs columns N must be positive (N = {N})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + if rng_seed is not None: + torch.manual_seed(rng_seed) + + lhs = torch.randn((M, K), dtype=torch.float32, device=device) + lhs = lhs.to(preferred_element_type) + + if trans_rhs: + rhs = torch.randn((G, N, K), dtype=torch.float32, device=device).permute( + 0, 2, 1 + ) + else: + rhs = torch.randn((G, K, N), dtype=torch.float32, device=device) + rhs = rhs.to(preferred_element_type) + + group_sizes = ( + gen_uniform_group_sizes(M, G, device=device) + if unif_group_sizes + else gen_group_sizes(M, G, device=device, rng_seed=None) + ) + + return lhs, rhs, group_sizes + + +def gen_gmm_output( + M: int, + N: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, +) -> Tensor: + assert M > 0, f"Number of out rows M must be positive (M = {M})." + assert N > 0, f"Number of out columns N must be positive (N = {N})." + + out = torch.empty((M, N), dtype=preferred_element_type, device=device) + + return out + + +def gen_gmm_tensors( + M: int, + K: int, + N: int, + G: int, + num_group_sizes: int, + device: torch.device | str = DEVICE, + input_type: torch.dtype = DTYPE, + output_type: torch.dtype = DTYPE, + trans_lhs: bool = False, + trans_rhs: bool = TRANS_RHS, + rng_seed: int | None = RNG_SEED, + unif_group_sizes: bool = False, + use_bias: bool = False, +) -> tuple[Tensor, Tensor, list[Tensor], Tensor, Tensor | None]: + lhs, rhs, group_sizes_0 = gen_gmm_input( + M, + K, + N, + G, + device=device, + preferred_element_type=input_type, + trans_rhs=trans_rhs, + rng_seed=rng_seed, + unif_group_sizes=unif_group_sizes, + ) + multiple_group_sizes = gen_multiple_group_sizes( + num_group_sizes, M, G, device=device, rng_seed=None, group_sizes_0=group_sizes_0 + ) + out = gen_gmm_output(M, N, device=device, preferred_element_type=output_type) + bias = None + if use_bias: + torch.manual_seed(rng_seed + 1000) # Different seed for bias + bias = torch.randn(G, N, dtype=input_type, device=device) + + return lhs, rhs, multiple_group_sizes, out, bias + + +# GMM helpers: get information from tensors. +# ------------------------------------------------------------------------------ + + +def get_gmm_shape( + lhs: Tensor, rhs: Tensor, group_sizes: Tensor +) -> tuple[int, int, int, int]: + assert lhs.dim() == 2, f"lhs must have 2 dimensions (it's {lhs.dim()})." + assert rhs.dim() == 3, f"rhs must have 3 dimensions (it's {rhs.dim()})." + assert ( + group_sizes.dim() == 1 + ), f"group_sizes must have 1 dimension (it's {group_sizes.dim()})." + + M, lhs_k = lhs.shape + rhs_g, rhs_k, N = rhs.shape + group_sizes_g = group_sizes.shape[0] + + assert ( + lhs_k == rhs_k + ), f"K dimension of lhs and rhs don't match (lhs = {lhs_k}, rhs = {rhs_k})." + K = lhs_k + assert ( + rhs_g == group_sizes_g + ), f"G dimension of rhs and group_sizes don't match (rhs = {rhs_g}, group_sizes = {group_sizes_g})." + G = rhs_g + + assert M > 0, f"M must be positive, it's {M}." + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}" + assert G > 0, f"G must be positive, it's {G}" + + return M, K, N, G + + +def get_gmm_output( + M: int, + N: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, +) -> Tensor: + assert M > 0, f"Number of out rows M must be positive (M = {M})." + assert N > 0, f"Number of out columns N must be positive (N = {N})." + + if existing_out is not None: + assert ( + existing_out.device == device + ), f"Existing output device and provided device don't match (existing = {existing_out.device}, provided = {device})." + assert ( + existing_out.dtype == preferred_element_type + ), f"Existing output type and preferred output type don't match (existing = {existing_out.dtype}, preferred = {preferred_element_type})." + assert existing_out.shape == ( + M, + N, + ), f"Existing output shape and GMM shape don't match (existing = {tuple(existing_out.shape)}, provided = {(M, N)})." + return existing_out + + return gen_gmm_output( + M, + N, + device=device, + preferred_element_type=preferred_element_type, + ) + + +def get_gmm_transposition(lhs: Tensor, rhs: Tensor, out: Tensor) -> tuple[bool, int]: + assert lhs.dim() == 2, f"lhs must have 2 dimensions (it's {lhs.dim()})." + assert rhs.dim() == 3, f"rhs must have 3 dimensions (it's {rhs.dim()})." + assert out.dim() == 2, f"out must have 2 dimensions (it's {out.dim()})." + + lhs_m, lhs_k = lhs.shape + G, rhs_k, rhs_n = rhs.shape + out_m, out_n = out.shape + + assert ( + lhs_m == out_m + ), f"M dimension of lhs and out don't match (lhs = {lhs_m}, rhs = {out_m})." + M = lhs_m + assert ( + lhs_k == rhs_k + ), f"K dimension of lhs and rhs don't match (lhs = {lhs_k}, rhs = {rhs_k})." + K = lhs_k + assert ( + rhs_n == out_n + ), f"N dimension of rhs and out don't match (lhs = {rhs_n}, rhs = {out_n})." + N = rhs_n + + assert M > 0, f"M must be positive, it's {M}." + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}" + assert G > 0, f"G must be positive, it's {G}" + + is_lhs_row_major = lhs.stride() == (K, 1) + assert is_lhs_row_major, "lhs must be row-major." + is_rhs_row_major = rhs.stride() == (K * N, N, 1) + is_rhs_col_major = rhs.stride() == (K * N, 1, K) + assert ( + is_rhs_row_major != is_rhs_col_major + ), "rhs must be row-major or column-major." + is_out_row_major = out.stride() == (N, 1) + assert is_out_row_major, "out must be row-major." + + # Get rhs leading dimension according to transposition configuration. + ld_rhs = N if is_rhs_row_major else K + + return is_rhs_col_major, ld_rhs + + +# TGMM helpers: tensor generation. +# ------------------------------------------------------------------------------ + + +def gen_tgmm_input( + M: int, + K: int, + N: int, + G: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, + trans_lhs: bool = TRANS_LHS, + rng_seed: int | None = RNG_SEED, + unif_group_sizes: bool = False, +) -> tuple[Tensor, Tensor, Tensor]: + assert K > 0, f"Number of lhs rows K must be positive (M = {K})." + assert M > 0, f"Number of lhs columns / rhs rows M must be positive (K = {M})." + assert N > 0, f"Number of rhs columns N must be positive (N = {N})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + if rng_seed is not None: + torch.manual_seed(rng_seed) + + if trans_lhs: + lhs = torch.randn((M, K), dtype=torch.float32, device=device).T + else: + lhs = torch.randn((K, M), dtype=torch.float32, device=device) + lhs = lhs.to(preferred_element_type) + + rhs = torch.randn((M, N), dtype=torch.float32, device=device) + rhs = rhs.to(preferred_element_type) + + group_sizes = ( + gen_uniform_group_sizes(M, G, device=device) + if unif_group_sizes + else gen_group_sizes(M, G, device=device, rng_seed=None) + ) + + return lhs, rhs, group_sizes + + +def gen_tgmm_output( + K: int, + N: int, + G: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, +) -> Tensor: + assert K > 0, f"Number of out rows K must be positive (K = {K})." + assert N > 0, f"Number of out columns N must be positive (N = {N})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + out = torch.empty((G, K, N), dtype=preferred_element_type, device=device) + + return out + + +def gen_tgmm_bias_grad( + K: int, + G: int, + device: torch.device | str = DEVICE, + with_bias_grad: bool = False, +) -> Tensor: + if with_bias_grad: + assert K > 0, f"Number of bias_grad rows K must be positive (K = {K})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + return torch.empty((G, K), device=device, dtype=torch.float32) + else: + # Return dummy pointer when bias_grad is not needed. + # Must be float32 because atomic_add does not support bf16/fp16, + # and Triton validates the pointer dtype even in dead branches. + return torch.tensor([], device=device, dtype=torch.float32) + + +def gen_tgmm_tensors( + M: int, + K: int, + N: int, + G: int, + num_group_sizes: int, + device: torch.device | str = DEVICE, + input_type: torch.dtype = DTYPE, + output_type: torch.dtype = DTYPE, + trans_lhs: bool = TRANS_LHS, + trans_rhs: bool = False, + rng_seed: int | None = RNG_SEED, + unif_group_sizes: bool = False, + use_bias: bool = False, +) -> tuple[Tensor, Tensor, list[Tensor], Tensor, Tensor | None]: + lhs, rhs, group_sizes_0 = gen_tgmm_input( + M, + K, + N, + G, + device=device, + preferred_element_type=input_type, + trans_lhs=trans_lhs, + rng_seed=rng_seed, + unif_group_sizes=unif_group_sizes, + ) + multiple_group_sizes = gen_multiple_group_sizes( + num_group_sizes, M, G, device=device, rng_seed=None, group_sizes_0=group_sizes_0 + ) + out = gen_tgmm_output(K, N, G, device=device, preferred_element_type=output_type) + if use_bias: + bias_grad = gen_tgmm_bias_grad(K, G, device=device, with_bias_grad=True) + else: + bias_grad = None + return lhs, rhs, multiple_group_sizes, out, bias_grad + + +# TGMM helpers: get information from tensors. +# ------------------------------------------------------------------------------ + + +def get_tgmm_shape( + lhs: Tensor, rhs: Tensor, group_sizes: Tensor +) -> tuple[int, int, int, int]: + assert lhs.dim() == 2, f"lhs must have 2 dimensions (it's {lhs.dim()})." + assert rhs.dim() == 2, f"rhs must have 2 dimensions (it's {rhs.dim()})." + assert ( + group_sizes.dim() == 1 + ), f"group_sizes must have 1 dimension (it's {group_sizes.dim()})." + + K, lhs_m = lhs.shape + rhs_m, N = rhs.shape + G = group_sizes.shape[0] + + assert ( + lhs_m == rhs_m + ), f"M dimension of lhs and rhs don't match (lhs = {lhs_m}, rhs = {rhs_m})." + M = lhs_m + + assert M > 0, f"M must be positive, it's {M}." + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}" + assert G > 0, f"G must be positive, it's {G}" + + return M, K, N, G + + +def get_tgmm_output( + K: int, + N: int, + G: int, + device: torch.device | str = DEVICE, + preferred_element_type: torch.dtype = DTYPE, + existing_out: Tensor | None = None, +) -> Tensor: + assert K > 0, f"Number of out rows K must be positive (K = {K})." + assert N > 0, f"Number of out columns N must be positive (N = {N})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + if existing_out is not None: + assert ( + existing_out.device == device + ), f"Existing output device and provided device don't match (existing = {existing_out.device}, provided = {device})." + assert ( + existing_out.dtype == preferred_element_type + ), f"Existing output type and preferred output type don't match (existing = {existing_out.dtype}, preferred = {preferred_element_type})." + assert existing_out.shape == ( + G, + K, + N, + ), f"Existing output shape and GMM shape don't match (existing = {tuple(existing_out.shape)}, provided = {(G, K, N)})." + return existing_out + + return gen_tgmm_output( + K, + N, + G, + device=device, + preferred_element_type=preferred_element_type, + ) + + +def get_tgmm_bias_grad( + K: int, + G: int, + device: torch.device | str = DEVICE, + existing_bias_grad: Tensor | None = None, +) -> Tensor: + """ + Get or validate bias gradient tensor for TGMM. + + If existing_bias_grad is provided, validates its shape, device, dtype, and stride, + and always zeros it before returning (since the kernel uses atomic_add). + If existing_bias_grad is None, returns a dummy tensor (for use when COMPUTE_BIAS_GRAD=False). + Parameters + ---------- + K : int + Number of rows in the bias gradient tensor. + G : int + Number of groups. + device : torch.device or str + Device for the tensor. + existing_bias_grad : torch.Tensor or None + Existing bias gradient tensor to validate and use. + Returns + ------- + torch.Tensor + Valid bias gradient tensor or dummy tensor. + """ + assert K > 0, f"Number of bias_grad rows K must be positive (K = {K})." + assert G > 0, f"Number of groups G must be positive (G = {G})." + + if existing_bias_grad is not None: + # Validate existing bias_grad tensor. + expected_shape = (G, K) + assert ( + tuple(existing_bias_grad.shape) == expected_shape + ), f"bias_grad must have shape {expected_shape}, got {tuple(existing_bias_grad.shape)}." + assert ( + existing_bias_grad.device == device + ), f"bias_grad must be on the same device (bias_grad = {existing_bias_grad.device}, device = {device})." + assert ( + existing_bias_grad.dtype == torch.float32 + ), f"bias_grad must be torch.float32 (kernel uses atomic_add which requires float32), got {existing_bias_grad.dtype}." + assert existing_bias_grad.stride() == ( + K, + 1, + ), f"bias_grad must be row-major with stride (K, 1) = ({K}, 1), got {existing_bias_grad.stride()}." + + # Always zero the tensor since bias_grad represents gradients for the current + # computation and should start fresh. The kernel uses atomic_add which adds to + # existing values, so we must zero before the kernel runs. + existing_bias_grad.zero_() + + return existing_bias_grad + + else: + return gen_tgmm_bias_grad(K, G, device=device, with_bias_grad=False) + + +def get_tgmm_transposition(lhs: Tensor, rhs: Tensor, out: Tensor) -> tuple[bool, int]: + assert lhs.dim() == 2, f"lhs must have 2 dimensions (it's {lhs.dim()})." + assert rhs.dim() == 2, f"rhs must have 2 dimensions (it's {rhs.dim()})." + assert out.dim() == 3, f"out must have 3 dimensions (it's {out.dim()})." + + lhs_k, lhs_m = lhs.shape + rhs_m, rhs_n = rhs.shape + G, out_k, out_n = out.shape + + assert ( + lhs_m == rhs_m + ), f"M dimension of lhs and rhs don't match (lhs = {lhs_m}, rhs = {rhs_m})." + M = lhs_m + assert ( + lhs_k == out_k + ), f"K dimension of lhs and out don't match (lhs = {lhs_k}, rhs = {out_k})." + K = lhs_k + assert ( + rhs_n == out_n + ), f"N dimension of rhs and out don't match (lhs = {rhs_n}, rhs = {out_n})." + N = rhs_n + + assert M > 0, f"M must be positive, it's {M}." + assert K > 0, f"K must be positive, it's {K}." + assert N > 0, f"N must be positive, it's {N}" + assert G > 0, f"G must be positive, it's {G}" + + is_lhs_row_major = lhs.stride() == (M, 1) + is_lhs_col_major = lhs.stride() == (1, K) + assert ( + is_lhs_row_major != is_lhs_col_major + ), "lhs must be row-major or column-major." + is_rhs_row_major = rhs.stride() == (N, 1) + assert is_rhs_row_major, "rhs must be row-major." + is_out_row_major = out.stride() == (K * N, N, 1) + assert is_out_row_major, "out must be row-major." + + # Get lhs leading dimension according to transposition configuration. + ld_lhs = M if is_lhs_row_major else K + + return is_lhs_col_major, ld_lhs diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/logger.py b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..391ddf9b6543f5244e7f4932c8568d60748e15cd --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_grouped_gemm_triton/utils/logger.py @@ -0,0 +1,47 @@ +import os +import logging + + +# AITER Triton Logger which is singleton object around python logging. +# Note: Python logging is also a singleton object, but we want to read the +# env var AITER_LOG_LEVEL once at the beginning. Another alternative is to do +# this in __init__.py. In fact, that's how CK logger is setup. We can look at +# switching to that at some point +# +# AITER_LOG_LEVEL follows python logging levels +# DEBUG +# INFO +# WARNING +# ERROR +# CRITICAL +# +class AiterTritonLogger(object): + _instance = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super(AiterTritonLogger, cls).__new__(cls) + log_level_str = os.getenv("AITER_TRITON_LOG_LEVEL", "WARNING").upper() + numeric_level = getattr(logging, log_level_str, logging.WARNING) + cls._instance._logger = logging.getLogger("AITER_TRITON") + cls._instance._logger.setLevel(numeric_level) + + return cls._instance + + def get_logger(self): + return self._logger + + def debug(self, msg): + self._logger.debug(msg) + + def info(self, msg): + self._logger.info(msg) + + def warning(self, msg): + self._logger.warning(msg) + + def error(self, msg): + self._logger.error(msg) + + def critical(self, msg): + self._logger.critical(msg) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a720e7a2cc4e44636f6e433a2750e945dc38e8b2 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/__init__.py @@ -0,0 +1,10 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +# from megablocks.layers.dmoe import dMoE +from .moe import MoE + +__all__ = [ + 'MoE', + # 'dMoE', +] diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/activation_fn.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/activation_fn.py new file mode 100644 index 0000000000000000000000000000000000000000..0e1d956704840aa4daf7d1d71d24e051567feab9 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/activation_fn.py @@ -0,0 +1,33 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Callable, Union + +import torch +from ..stk import Matrix + + +def act_fn( + x: Matrix, + function: Callable, + return_grad_fn: bool = False, + **kwargs, +) -> Union[tuple[Matrix, Any] | Matrix]: + assert isinstance(x, Matrix) + with torch.set_grad_enabled(torch.is_grad_enabled() or return_grad_fn): + if return_grad_fn: + x.data.requires_grad = True + out = function(x.data, **kwargs) + y = Matrix( + x.size(), + out, + x.row_indices, + x.column_indices, + x.offsets, + x.column_indices_t, + x.offsets_t, + x.block_offsets_t, + ) + if return_grad_fn: + return y, out.backward + return y diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/all_to_all.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/all_to_all.py new file mode 100644 index 0000000000000000000000000000000000000000..5ac7067bcaa34db1d82b340c43550fe3577aa7a3 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/all_to_all.py @@ -0,0 +1,54 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.distributed as dist + + +class AllToAllOp(torch.autograd.Function): + + @staticmethod + def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op): + out = torch.empty((sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype) + + ctx.input_shape = x.shape + ctx.output_split_sizes = output_split_sizes + ctx.input_split_sizes = input_split_sizes + ctx.group = group + handle = dist.all_to_all_single( + out, + x, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group, + async_op=async_op, + ) + return out, handle + + @staticmethod + def backward(ctx, grad, _): + if ctx.needs_input_grad[0]: + out = torch.empty( + ctx.input_shape, + device=grad.device, + dtype=grad.dtype, + ) + dist.all_to_all_single( + out, + grad, + output_split_sizes=ctx.input_split_sizes, + input_split_sizes=ctx.output_split_sizes, + group=ctx.group, + ) + return out, None, None, None, None + return None, None, None, None, None + + +def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False): + return AllToAllOp.apply( + x, + output_split_sizes, + input_split_sizes, + group, + async_op, + ) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/arguments.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/arguments.py new file mode 100644 index 0000000000000000000000000000000000000000..4db9b1bd38bc2e2f421625c124f86b85f45c5ae0 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/arguments.py @@ -0,0 +1,101 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +from functools import partial +from typing import Any, Callable, Optional, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +# import megablocks.grouped_gemm_util as grouped_gemm +from .. import grouped_gemm_util as grouped_gemm + +# Type annotation for in-place Tensor initialization function. +InitFn = Union[Callable[[torch.Tensor], None], partial[torch.Tensor]] + +_ALLOWED_BITWIDTHS = (-1, 4, 8) + +DEFAULT_ACTIVATION_FN = partial(F.gelu, approximate='tanh') + + +@dataclasses.dataclass +class Arguments: + # Model arguments. + hidden_size: int = 1024 + ffn_hidden_size: int = 4096 + num_layers: int = 1 + bias: bool = True + return_bias: bool = True + activation_fn: Optional[Callable] = DEFAULT_ACTIVATION_FN + + # MoE arguments. + moe_num_experts: int = 1 + moe_top_k: int = 1 + moe_capacity_factor: int = 1 + moe_normalize_expert_weights: Optional[Union[int, float]] = None + moe_loss_weight: float = 0.1 + moe_jitter_eps: Optional[float] = None + moe_lbl_in_fp32: bool = False + + # Parallelism arguments. + moe_expert_model_parallelism: bool = False + expert_parallel_group: Optional[dist.ProcessGroup] = None + pipeline_model_parallel_size: int = 1 + num_layers_per_virtual_pipeline_stage: Optional[int] = None + + # Compute arguments. + memory_optimized_mlp: bool = False + mlp_type: str = 'mlp' + mlp_impl: str = 'sparse' + + # Initialization arguments. + fp16: bool = True + bf16: bool = False + device: Union[int, torch.device] = dataclasses.field(default_factory=torch.cuda.current_device) + init_method: InitFn = partial(torch.nn.init.normal_, mean=0.0, std=0.02) + output_layer_init_method: InitFn = init_method + + # Benchmarking arguments. + uniform_expert_assignment: bool = False + + # shared expert arguments + shared_expert: bool = False # enable using shared expert + fc_cls: Any = torch.nn.Linear # class of the fully connected layer in shared expert (purpose: to allow using custom FC layer eg te.Linear (for FP8)) + fc_kwargs: dict[str, Any] = dataclasses.field(default_factory=dict,) # kwargs for custom fc layers + remat_act_fn: bool = True # enable act fn to be rematerialized instead of stored + shared_expert_hidden_size: Optional[ + int] = None # hidden size of the shared expert IF we want to set it to something different from hidden_size + shared_expert_weighted_sum: bool = False # enable using weighted sum for shared expert output (wieghted by number of experts used) + + # Router Z-loss arguments + moe_zloss_weight: float = 0 # 1e-3 is a reasonable value + moe_zloss_in_fp32: bool = False + + def __post_init__(self): + # Sparse MLP is not supported with triton >=3.2.0 + # TODO: Remove this once sparse is supported with triton >=3.2.0 + if self.__getattribute__('mlp_impl') == 'sparse': + try: + import triton + if triton.__version__ >= '3.2.0': + raise ValueError( + 'Sparse MLP is not supported with triton >=3.2.0. Please use mlp_impl="grouped" instead.', + ) + except ImportError: + raise ImportError('Triton is required for sparse MLP implementation') + + if self.__getattribute__('mlp_impl') == 'grouped': + grouped_gemm.assert_grouped_gemm_is_available() + + if self.shared_expert_hidden_size is None: + self.shared_expert_hidden_size = self.ffn_hidden_size + + +def from_megatron(megatron_args: Any): + args = Arguments() + for field in dataclasses.fields(args): + if hasattr(megatron_args, field.name): + setattr(args, field.name, getattr(megatron_args, field.name)) + return args diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/common.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/common.py new file mode 100644 index 0000000000000000000000000000000000000000..2d07109702963ba48a3b94ab860807954dfd79c1 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/common.py @@ -0,0 +1,26 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from .arguments import Arguments + + +def dtype(args: Arguments): + if args.fp16: + return torch.float16 + elif args.bf16: + return torch.bfloat16 + return None + + +def cast_if_autocast_enabled(tensor): + if torch.is_autocast_enabled(): + if tensor.device.type == 'cuda': + dtype = torch.get_autocast_gpu_dtype() + elif tensor.device.type == 'cpu': + dtype = torch.get_autocast_cpu_dtype() + else: + raise NotImplementedError() + return tensor.to(dtype=dtype) + return tensor diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/dmlp_registry.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/dmlp_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..de2ed047042e438c7190ebb139b6f7f30009734c --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/dmlp_registry.py @@ -0,0 +1,42 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Union + +from . import glu, mlp +from .arguments import Arguments + +MlpType = Union[mlp.SparseMLP, glu.SparseGLU] + +_REGISTRY = { + 'mlp': { + 'grouped': mlp.GroupedMLP, + 'sparse': mlp.SparseMLP, + }, + 'glu': { + 'grouped': glu.GroupedGLU, + 'sparse': glu.SparseGLU, + }, +} + + +def get(args: Arguments) -> MlpType: + """Returns an MLP for use in a dMoE instance. + + Uses the provided arguments to instantiate the appropriate + MLP instance. This only contains MLPs for use in dMoEs + (ie. only for the dropless versions of MoEs). + + Args: + args: propagated Arguments dataclass. + + Returns: + An instantiated MLP constructed using the input args. + """ + if args.mlp_type not in _REGISTRY: + raise ValueError(f'Unsupported mlp type: {args.mlp_type}') + + if args.mlp_impl not in _REGISTRY[args.mlp_type]: + raise ValueError(f'{args.mlp_type} does not support {args.mlp_impl} backend.',) + + return _REGISTRY[args.mlp_type][args.mlp_impl](args) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/dmoe.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/dmoe.py new file mode 100644 index 0000000000000000000000000000000000000000..6d0375a4df2f27134c4127e60be04f3b45693050 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/dmoe.py @@ -0,0 +1,337 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import torch + +# try: +# import stk.ops +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/_layers/dmoe.py is needed.', +# ) + +# import megablocks.ops as ops +# # from megablocks.ops import ops +# from megablocks.layers import common, dmlp_registry, moe, mpu +# from megablocks.layers.arguments import Arguments + +from .. import stk +from .. import ops +from . import common, dmlp_registry, moe, mpu +from .arguments import Arguments + +def promote_scalar(x): + return x.view(1) if not len(x.size()) else x + + +class ParallelDroplessMLP(moe.ParallelMLP): + + def __init__(self, args: Arguments): + super(ParallelDroplessMLP, self).__init__(args) + self.hidden_size = args.hidden_size + self.ffn_hidden_size = mpu.features_per_rank(args) + self.blocking = 128 + self.mlp = dmlp_registry.get(args) + + # Calculate the number of bits needed to represent the column indices + # in the intermediate sparse matrix. + max_column_index = ((self.ffn_hidden_size * self.num_experts) // self.blocking) + self.transpose_sort_end_bit = max( + int(np.ceil(np.log2(max_column_index))), + 1, + ) + + def sparse_transpose(self, size, row_indices, column_indices, offsets): + block_columns = size[1] // self.blocking + + # Sort row indices by column indices to get the transposed matrix's + # column indices. + # + # NOTE: Our sort operation uses the same width indices as the input values. + # To avoid overflow when we have large activation matrices we cast to + # 32-bit before sorting. + _, gather_indices = ops.sort( + column_indices.int(), + self.transpose_sort_end_bit, + ) + + # There are a constant number of blocks in every row of the sparse matrix. + # A blocks offset is: + # + # row_index * blocks_per_row + column_index % blocks_per_row + # + # Once we have the block offsets ordered for transposition we can divide + # by blocks_per_row to get the transposed column indices. + column_indices_t = row_indices.gather(0, gather_indices.long()) + block_offsets_t = gather_indices.int() + + zero = torch.zeros((1,), dtype=torch.int32, device=row_indices.device) + nnz_per_column = ops.histogram(column_indices, block_columns) + nnz_per_column = ops.inclusive_cumsum(nnz_per_column, 0) + if nnz_per_column.dim() == 0: + # This addresses an edge case when ffn_hidden_size is equal to self.blocking. + nnz_per_column = nnz_per_column.unsqueeze(0) + offsets_t = torch.cat([zero, nnz_per_column]) + return column_indices_t, offsets_t, block_offsets_t + + def topology(self, x, padded_bins): + padded_tokens, _ = x.size() + assert padded_tokens % self.blocking == 0 + if self.ffn_hidden_size % self.blocking != 0: + raise ValueError( + f'The ffn_hidden_size {self.ffn_hidden_size} must be divisible by ' + + f'the block size {self.blocking}. Please update your configuration.', + ) + + # Offsets for the sparse matrix. All rows have the + # same number of nonzero blocks dictated by the + # dimensionality of a single expert. + block_rows = padded_tokens // self.blocking + blocks_per_row = self.ffn_hidden_size // self.blocking + offsets = torch.arange( + 0, + block_rows * blocks_per_row + 1, + blocks_per_row, + dtype=torch.int32, + device=x.device, + ) + + # Indices for the sparse matrix. The indices for + # the intermediate matrix are dynamic depending + # on the mapping of tokens to experts. + column_indices = ops.topology( + padded_bins, + self.blocking, + block_rows, + blocks_per_row, + ) + + # TODO(tgale): This is unused. Remove the need for this in stk. + # For now, use meta init to save the device memory. + data = torch.empty( + column_indices.numel(), + self.blocking, + self.blocking, + dtype=common.dtype(self.args), + device='meta', + ) + shape = ( + padded_tokens, + self.ffn_hidden_size * mpu.experts_per_rank(self.args), + ) + row_indices = stk.ops.row_indices(shape, data, offsets, column_indices) + column_indices_t, offsets_t, block_offsets_t = self.sparse_transpose( + shape, + row_indices, + column_indices, + offsets, + ) + return stk.Matrix( + shape, + data, + row_indices, + column_indices, + offsets, + column_indices_t, + offsets_t, + block_offsets_t, + ) + + def indices_and_padded_bins(self, top_experts): + # Sort the expert ids to produce the scatter/gather + # indices for the permutation. + top_experts = top_experts.int() + bin_ids, indices = ops.sort(top_experts, self.sort_end_bit) + + # Histogram the expert ids to identify the number of + # tokens routed to each expert. + tokens_per_expert = ops.histogram(top_experts, self.num_experts) + + # Round the token counts up to the block size used in + # the matrix muliplications. Caculate the starting + # position of each bin. + padded_tokens_per_expert = ops.round_up( + tokens_per_expert, + self.blocking, + ) + padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) + padded_bins = promote_scalar(padded_bins) + + # Calculate the bin bounds for the sorted tokens. + bins = ops.inclusive_cumsum(tokens_per_expert, 0) + bins = promote_scalar(bins) + return indices, bin_ids, bins, padded_bins, tokens_per_expert + + def sparse_forward_once(self, x, expert_weights, top_experts): + # x: [sl, bs, hs] + # expert_weights: [sl * bs, top-k] + # top_experts: [sl * bs, top-k] + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + with torch.no_grad(): + indices, bin_ids, bins, padded_bins, tokens_per_expert = (self.indices_and_padded_bins(top_experts)) + + # Route the tokens for MoE computation. + x = x.view(-1, x.shape[-1]) + x = ops.padded_gather( + x, + indices, + bin_ids, + bins, + padded_bins, + self.top_k, + ) + + # Create the sparse matrix topology. + with torch.no_grad(): + topo = self.topology(x, padded_bins) + + # Perform the expert computation. + x = self.mlp(x, topo) + + # Un-route the data for the MoE output. + x = ops.padded_scatter( + x, + indices, + bin_ids, + expert_weights, + bins, + padded_bins, + self.top_k, + ) + return x, tokens_per_expert + + # For use in the base-class parallel_forward_once. + def sparse_permute_and_compute( + self, + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, # unused + top_k, + ): + + # Round the token counts up to the block size used in the matrix + # multiplication. Calculate the starting position of each bin. + padded_tokens_per_expert = ops.round_up( + tokens_per_expert, + self.blocking, + ) + padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) + padded_bins = promote_scalar(padded_bins) + + # Route the tokens for MoE computation. + x = x.view(-1, x.shape[-1]) + x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k) + + # Create the sparse matrix topology. + with torch.no_grad(): + topo = self.topology(x, padded_bins) + + # Perform the expert computation. + x = self.mlp(x, topo) + + # Un-route the data for the MoE output. + return ops.padded_scatter( + x, + indices, + bin_ids, + expert_weights, + bins, + padded_bins, + top_k, + ) + + def grouped_forward_once(self, x, expert_weights, top_experts): + # x: [sl, bs, hs] + # expert_weights: [sl * bs, top-k] + # top_experts: [sl * bs, top-k] + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + with torch.no_grad(): + indices, bin_ids, bins, tokens_per_expert = (self.indices_and_bins(top_experts)) + + out = self.grouped_permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + -1, # unused + self.args.moe_top_k, + ) + return out, tokens_per_expert + + def grouped_permute_and_compute( + self, + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, # unused + top_k, + ): + + # Route the tokens for MoE computation. + x = x.view(-1, x.shape[-1]) + x = ops.gather(x, indices, bin_ids, bins, top_k) + + # Perform the expert computation. + x = self.mlp(x, tokens_per_expert) + + # Un-route the data for the MoE output. + return ops.scatter(x, indices, bin_ids, expert_weights, bins, top_k) + + def forward_once(self, x, expert_weights, top_experts): + if self.args.mlp_impl == 'sparse': + return self.sparse_forward_once(x, expert_weights, top_experts) + else: + return self.grouped_forward_once(x, expert_weights, top_experts) + + def permute_and_compute( + self, + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, + top_k, + ): + if self.args.mlp_impl == 'sparse': + return self.sparse_permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, + top_k, + ) + else: + return self.grouped_permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capactiy, + top_k, + ) + + +class dMoE(moe.MoE): + + def _init_experts_mlp(self, args: Arguments): + return ParallelDroplessMLP(args) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/gelu.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/gelu.py new file mode 100644 index 0000000000000000000000000000000000000000..c4c9e6532798615b5c12c96694241a4c18ee8f7b --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/gelu.py @@ -0,0 +1,52 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +# try: +# import stk +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/_layers/gelu.py is needed.', +# ) + +from .. import stk + +import torch +import torch.nn.functional as F + + +@torch.jit.script +def _gelu_backward_inplace(g, x): + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + ff = (0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)) + return g.mul_(ff) + + +def gelu_backward_(grad: stk.Matrix, x: stk.Matrix): + # NOTE: The two sparse matrices must have the same topology. + if isinstance(grad, stk.Matrix) and isinstance(x, stk.Matrix): + return stk.Matrix( + x.size(), + _gelu_backward_inplace(grad.data, x.data), + x.row_indices, + x.column_indices, + x.offsets, + x.column_indices_t, + x.offsets_t, + x.block_offsets_t, + ) + return _gelu_backward_inplace(grad, x) + + +def gelu(x: stk.Matrix): + assert isinstance(x, stk.Matrix) + return stk.Matrix( + x.size(), + F.gelu(x.data, approximate='tanh'), + x.row_indices, + x.column_indices, + x.offsets, + x.column_indices_t, + x.offsets_t, + x.block_offsets_t, + ) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/glu.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/glu.py new file mode 100644 index 0000000000000000000000000000000000000000..5f297a41ff6a1a2a285f5b461951672364b898da --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/glu.py @@ -0,0 +1,244 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +# import stk.ops +# try: +# import stk.ops +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/_layers/glu.py is needed.', +# ) + +from .. import stk + +import torch + +# from megablocks import grouped_gemm_util as gg +# from megablocks.layers import common, mpu +# from megablocks.layers.activation_fn import act_fn +# from megablocks.layers.arguments import Arguments +# from megablocks.layers.mlp import ( +# SharedMLP, +# SparseMLP, +# create_dmoe_expert_weights, +# resolve_dtensor, +# ) + +from .. import grouped_gemm_util as gg +from . import common, mpu +from .activation_fn import act_fn +from .arguments import Arguments +from .mlp import ( + SharedMLP, + SparseMLP, + create_dmoe_expert_weights, + resolve_dtensor, +) + + +class SparseGLU(SparseMLP): + + def __init__(self, args: Arguments): + super().__init__(args) + self.v1 = torch.nn.Parameter( + torch.empty( + self._num_rows_per_rank, + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + with torch.no_grad(): + self.v1.copy_( + create_dmoe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.init_method, + ), + ) + + mpu.set_expert_model_parallel_attributes( + self.v1, + self._should_set_parallelism_attribute, + ) + + def forward(self, x, topo): + if self.args.memory_optimized_mlp: + raise NotImplementedError( + 'Memory optimized implementation not yet supported with GLU with sparse kernels.', + ) + + w1, v1, w2 = self.scale_grad(self.w1), self.scale_grad(self.v1,), self.scale_grad(self.w2) + w1, v1, w2 = resolve_dtensor(w1), resolve_dtensor(v1,), resolve_dtensor(w2) + + # Compute the GLU. + x1 = stk.ops.sdd(x, w1.t(), topo) + x2 = stk.ops.sdd(x, v1.t(), topo) + + activation_fn_out = act_fn(x1, self.args.activation_fn) + x1 = stk.ops.mul(activation_fn_out, x2) + + return stk.ops.dsd(x1, w2) + + +class MemoryOptimizedGroupedGLU(torch.autograd.Function): + """GroupedMLP with manually scheduled memory reuse.""" + + @staticmethod + @torch.amp.autocast_mode.custom_fwd(device_type='cuda') + def forward(ctx, x, w1, v1, w2, batch_sizes, activation_fn): + # Cast inputs using ctx dtype from AMP + if ctx._fwd_used_autocast: + x = x.to(ctx._dtype) + w1 = w1.to(ctx._dtype) + v1 = v1.to(ctx._dtype) + w2 = w2.to(ctx._dtype) + # x: [m, k], w1: [n, k], v1: [n, k], w2: [n, k] + if (not x.is_contiguous() or not w1.is_contiguous() or not v1.is_contiguous() or not w2.is_contiguous()): + raise ValueError("Expected contiguous 'x', 'w1', 'v1' and 'w2'.") + + # Layer 0: x @ w1.t(). + assert gg.backend is not None + sdd_out = gg.backend.gmm(x, w1, batch_sizes, trans_b=True) + v1_out = gg.backend.gmm(x, v1, batch_sizes, trans_b=True) + + # GeLU. + activation_fn_out = activation_fn(sdd_out) * v1_out + + # Layer 1: x @ w2. + dsd_out = gg.backend.gmm(activation_fn_out, w2, batch_sizes) + + # NOTE: Save the input to the layer and the activation_fn input for + # gradient computation. We'll re-compute the activation_fn forward + # pass in the backward pass to avoid materializing another + # intermediate. + ctx.x_shape = x.shape + ctx.sdd_out_shape = sdd_out.shape + ctx.dtype = x.dtype + ctx.activation_fn = activation_fn + ctx.save_for_backward(w1, v1, w2, batch_sizes, x, sdd_out, v1_out) + return dsd_out + + @staticmethod + @torch.amp.autocast_mode.custom_bwd(device_type='cuda') + def backward(ctx, ddsd_out): + if (not ctx.needs_input_grad[0] or not ctx.needs_input_grad[1] or not ctx.needs_input_grad[2]): + raise ValueError('Expected all MLP inputs to need grad.') + + # Unpack saved tensors + # dtype = ctx.dtype + saved_tensors = ctx.saved_tensors + w1, v1, w2 = saved_tensors[:3] + batch_sizes = saved_tensors[3] + x = saved_tensors[4] + sdd_out, v1_out = saved_tensors[5:7] + + # Rematerialize activation_fn output. + activation_fn = ctx.activation_fn + with torch.set_grad_enabled(True): + sdd_out.requires_grad = True + v1_out.requires_grad = True + activation_fn_out = activation_fn(sdd_out) * v1_out + activation_grad_fn = activation_fn_out.backward + + # Compute dw2 with recomputed activation_fn output. + assert gg.backend is not None + dw2 = gg.backend.gmm( + activation_fn_out, + ddsd_out, + batch_sizes, + trans_a=True, + ) + + # Compute dactivation_fn_out. + # + # NOTE: We reuse the activation_fn_out allocation. + dactivation_fn_out = activation_fn_out + gg.backend.gmm( + ddsd_out, + w2, + batch_sizes, + trans_b=True, + c=dactivation_fn_out, + ) + + # Compute dsdd_out. + # + # NOTE: This reuses the dactivation_fn_out allocation. + assert activation_grad_fn is not None + activation_grad_fn(dactivation_fn_out) + dsdd_out = sdd_out.grad + dv1_out = v1_out.grad + + # Compute dw1. + dw1 = gg.backend.gmm(dsdd_out, x, batch_sizes, trans_a=True) + + # Compute dv1. + dv1 = gg.backend.gmm(dv1_out, x, batch_sizes, trans_a=True) + + # Compute dx. + # + # NOTE: This reuses the ddsd_out allocation. + dx = ddsd_out + gg.backend.gmm(dsdd_out, w1, batch_sizes, c=dx) + dx += gg.backend.gmm(dv1_out, v1, batch_sizes) + return dx, dw1, dv1, dw2, None, None + + +memory_optimized_grouped_glu = MemoryOptimizedGroupedGLU.apply + + +class GroupedGLU(SparseGLU): + + def forward(self, x, tokens_per_expert): + batch_sizes = tokens_per_expert.cpu().to(torch.long) + w1, v1, w2 = ( + self.scale_grad(self.w1), + self.scale_grad(self.v1), + self.scale_grad(self.w2), + ) + w1, v1, w2 = resolve_dtensor(w1), resolve_dtensor(v1,), resolve_dtensor(w2) + + # Re-shape the weights for the grouped GEMMs. + ne = mpu.experts_per_rank(self.args) + w1 = w1.view(ne, -1, self.args.hidden_size) + v1 = v1.view(ne, -1, self.args.hidden_size) + w2 = w2.view(ne, -1, self.args.hidden_size) + + if self.args.memory_optimized_mlp: + return memory_optimized_grouped_glu( + x, + w1, + v1, + w2, + batch_sizes, + self.args.activation_fn, + ) + + # Compute the MLP. + assert gg.ops is not None + x1 = gg.ops.gmm(x, w1, batch_sizes, trans_b=True) + x2 = gg.ops.gmm(x, v1, batch_sizes, trans_b=True) + x1 = self.args.activation_fn(x1) * x2 + return gg.ops.gmm(x1, w2, batch_sizes) + + +class SharedGLU(SharedMLP): + """GPU for shared expert. + + Note: this is a copy -> pasta -> modify of the LLM-Foundry MPTGLU class + """ + + def __init__(self, args: Arguments): + super().__init__(args) + self.gate_proj = args.fc_cls( + args.hidden_size, + self.args.shared_expert_hidden_size, + **self.fc_kwargs, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x)) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/memory_test.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/memory_test.py new file mode 100644 index 0000000000000000000000000000000000000000..74d1166931b712635131985b25a89f4ca23e576d --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/memory_test.py @@ -0,0 +1,103 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import gc + +import torch +import torch.distributed as dist + +# from megablocks.layers import arguments, dmoe +from . import arguments, dmoe + +_TESTS = ((8, 2048, 4096, 4096, 32, 4),) + + +def get_tensors(): + ptrs = set() + out = [] + for obj in gc.get_objects(): + if torch.is_tensor(obj): + if not obj.is_contiguous() or obj.data_ptr() in ptrs: + continue + out.append(obj) + ptrs.add(obj.data_ptr()) + return out + + +def test_memory( + group, + batch_size, + sequence_length, + hidden_size, + ffn_hidden_size, + num_experts, + top_k, +): + args = arguments.Arguments( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + moe_num_experts=num_experts, + moe_top_k=top_k, + moe_expert_model_parallelism=True, + expert_parallel_group=group, + fp16=False, + bf16=True, + device=torch.cuda.current_device(), + ) + layer = dmoe.dMoE(args).cuda() + + x = torch.randn((batch_size, sequence_length, hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16).requires_grad_(True) + torch.cuda.empty_cache() + + # Run forward + backward. + # with torch.autograd.detect_anomaly(): + out, _ = layer(x) + out.mean().backward() + + # Report peak memory. + mem = torch.cuda.max_memory_allocated() + print('Max Memory Allocated = {:0.0f}MiB'.format(mem / 1e6)) + print('Max Memory Reserved = {:0.0f}MiB'.format(torch.cuda.max_memory_reserved() / 1e6,),) + + # Calculate weight and gradient memory usage. + weight_memory = 2 * ( + layer.router.layer.weight.numel() + layer.experts.mlp.w1.numel() + layer.experts.mlp.w2.numel() + ) + + def grad_numel(x): + if x.grad is not None: + return x.grad.numel() + return 0 + + grad_memory = 2 * ( + grad_numel(layer.router.layer.weight) + grad_numel(layer.experts.mlp.w1) + grad_numel(layer.experts.mlp.w2) + ) + weight_memory += grad_memory + + print('Weight Memory Allocated = {:0.0f}MiB'.format(weight_memory / 1e6)) + print('Activation Memory Allocated = {:0.0f}MiB'.format((mem - weight_memory) / 1e6,),) + + # Manually calculate GPU memory usage from the garbage + # collector. + gc.collect() + total = 0 + tensors = get_tensors() + tensors = sorted(tensors, key=lambda x: -x.numel()) + for i, t in enumerate(tensors): + total += t.numel() + print(f'{i}: {t.shape}, {t.numel() * 2}') + del tensors + + print('Total Bytes Found = {:0.0f}MiB'.format(total * 2 / 1e6)) + + +if __name__ == '__main__': + assert dist.is_available() + group = dist.init_process_group(backend='nccl') + local_rank = dist.get_rank(group) + torch.cuda.set_device(local_rank) + + for args in _TESTS: + test_memory(group, *args) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/mlp.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..c99afb9904c24a8b6a83e79059cd1251dbbfd99e --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/mlp.py @@ -0,0 +1,587 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# try: +# import stk +# import stk.backend.triton_kernels +# import stk.ops +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/_layers/mlp.py is needed.', +# ) + +from .. import stk + +import torch +from packaging import version + +# from megablocks import grouped_gemm_util as gg +# from megablocks.layers import common, gelu, mpu +# from megablocks.layers.activation_fn import act_fn +# from megablocks.layers.arguments import DEFAULT_ACTIVATION_FN, Arguments, InitFn + +from .. import grouped_gemm_util as gg +from . import common, gelu, mpu +from .activation_fn import act_fn +from .arguments import DEFAULT_ACTIVATION_FN, Arguments, InitFn + +class ScaleGradient(torch.autograd.Function): + + @staticmethod + @torch.amp.autocast_mode.custom_fwd(device_type='cuda') + def forward(ctx: Any, x: torch.Tensor, scale: float): + ctx.scale = scale + return x + + @staticmethod + @torch.amp.autocast_mode.custom_bwd(device_type='cuda') + def backward(ctx: torch.Tensor, grad: torch.Tensor): + return grad * ctx.scale, None + + +scale_gradient = ScaleGradient.apply + + +def resolve_dtensor(weight: torch.Tensor): + if version.parse(torch.__version__) >= version.parse('2.0.0'): + from torch.distributed._tensor import DTensor + if isinstance(weight, DTensor): + return weight.to_local() + return weight + + +def create_moe_expert_weights( + args: Arguments, + num_experts: int, + ffn_hidden_size: int, + hidden_size: int, + init_method: InitFn, +): + # Create the entire weight matrix such that the sampled weights will + # not vary between data parallelism and expert model parallelism for + # the same random seed. + master_weights = torch.empty( + num_experts, + ffn_hidden_size, + hidden_size, + device=args.device, + dtype=common.dtype(args), + ) + init_method(master_weights) + + if not args.moe_expert_model_parallelism: + return master_weights + + # Calculate the amount of sharding in each dimension. + expert_sharding_degree = mpu.expert_sharding_degree(args) + hidden_sharding_degree = mpu.hidden_sharding_degree(args) + + # Calculate the experts per rank. + # + # NOTE: We assign ranks to be expert parallel before going + # tensor parallel. + rank = mpu.get_expert_parallel_rank(args) + expert_rank = rank % expert_sharding_degree + num_experts_per_rank = num_experts // expert_sharding_degree + start_expert = expert_rank * num_experts_per_rank + end_expert = (expert_rank + 1) * num_experts_per_rank + + # Calculate the rows per rank. + row_rank = rank // expert_sharding_degree + num_rows_per_rank = ffn_hidden_size // hidden_sharding_degree + start_row = row_rank * num_rows_per_rank + end_row = (row_rank + 1) * num_rows_per_rank + + # Slice the weight matrix to get the chunk for this rank. + with torch.no_grad(): + weights = master_weights[start_expert:end_expert, start_row:end_row] + return weights + + +class MLP(torch.nn.Module): + + def __init__(self, args: Arguments): + super().__init__() + self.args = args + # expert_parallel_world_size = mpu.get_expert_parallel_world_size(args) + experts_per_rank = mpu.experts_per_rank(args) + + self.w1 = torch.nn.Parameter( + torch.empty( + experts_per_rank, + args.hidden_size, + mpu.features_per_rank(args), + device=args.device, + dtype=common.dtype(args), + ), + ) + self.w2 = torch.nn.Parameter( + torch.empty( + experts_per_rank, + mpu.features_per_rank(args), + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + mpu.set_expert_model_parallel_attributes( + self.w1, + args.moe_expert_model_parallelism, + ) + mpu.set_expert_model_parallel_attributes( + self.w2, + args.moe_expert_model_parallelism, + ) + + # Initialize the parameters for the MLP. + # + # NOTE: It is important that we create the weight tensors prior + # to creating the master weights and slicing our the piece for + # this rank. If the master weights are created first the PyTorch + # caching allocator appears to use the same memory block for these + # and the slice which causes large increases in our peak memory + # usage. + with torch.no_grad(): + w1 = create_moe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.init_method, + ) + self.w1.copy_(w1.transpose(1, 2).contiguous()) + self.w2.copy_( + create_moe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.output_layer_init_method, + ), + ) + + self.gradient_scale = None + if self.args.moe_expert_model_parallelism: + self.gradient_scale = 1 / mpu.get_expert_parallel_world_size(self.args,) + + def scale_grad(self, w): + if self.gradient_scale is None: + return w + return scale_gradient(w, self.gradient_scale) + + def forward(self, x): + w1, w2 = self.scale_grad(self.w1), self.scale_grad(self.w2) + w1, w2 = resolve_dtensor(w1), resolve_dtensor(w2) + x = torch.bmm(x, w1) + x = self.args.activation_fn(x) + return torch.bmm(x, w2) + + +def create_dmoe_expert_weights( + args: Arguments, + num_experts: int, + rows: int, + columns: int, + init_method: InitFn, +): + weights = create_moe_expert_weights( + args, + num_experts, + rows, + columns, + init_method, + ) + return weights.view([-1, columns]) + + +class MemoryOptimizedMLP(torch.autograd.Function): + """Sparse MLP with manually scheduled memory reuse.""" + + @staticmethod + @torch.amp.autocast_mode.custom_fwd(device_type='cuda') + def forward(ctx, x, w1, w2, topo, activation_fn): + # Cast inputs using ctx dtype from AMP + if ctx._fwd_used_autocast: + x = x.to(ctx._dtype) + w1 = w1.to(ctx._dtype) + w2 = w2.to(ctx._dtype) + # x: [m, k], w1: [n, k], w2: [n, k] + if (not x.is_contiguous() or not w1.is_contiguous() or not w2.is_contiguous()): + raise ValueError("Expected contiguous 'x', 'w1' and 'w2'.") + + topo_tensors = ( + topo.row_indices, + topo.column_indices, + topo.offsets, + topo.column_indices_t, + topo.offsets_t, + topo.block_offsets_t, + ) + + # Layer 0: x @ w1.t(). + sdd_out = stk.ops.sdd(x, w1.t(), topo) + + # GeLU. + activation_fn_out = act_fn(sdd_out, activation_fn) + + # Layer 1: x @ w2. + dsd_out = stk.ops.dsd(activation_fn_out, w2) + + # NOTE: Save the input to the layer and the activation_fn input for + # gradient computation. We'll re-compute the activation_fn forward + # pass in the backward pass to avoid materializing another + # intermediate. + ctx.shape = topo.shape + ctx.x_shape = x.shape + ctx.sdd_out_shape = sdd_out.data.shape + ctx.dtype = x.dtype + ctx.activation_fn = activation_fn + ctx.save_for_backward(w1, w2, *topo_tensors, x, sdd_out.data) + return dsd_out + + @staticmethod + @torch.amp.autocast_mode.custom_bwd(device_type='cuda') + def backward(ctx, ddsd_out): + if (not ctx.needs_input_grad[0] or not ctx.needs_input_grad[1] or not ctx.needs_input_grad[2]): + raise ValueError('Expected all MLP inputs to need grad.') + + # unpack saved tensors + # dtype = ctx.dtype + saved_tensors = ctx.saved_tensors + w1, w2 = saved_tensors[:2] + topo_tensors = saved_tensors[2:8] + x = saved_tensors[8] + sdd_out_data = saved_tensors[9] + + # rematerialize activation function output + activation_fn = ctx.activation_fn + sdd_out = stk.Matrix(ctx.shape, sdd_out_data, *topo_tensors) + activation_fn_out, activation_grad_fn = act_fn( + sdd_out, + activation_fn, + return_grad_fn=True, + ) + + # Compute dw2 with recomputed activation_fn output. + dw2 = stk.ops.dsd(activation_fn_out.t(), ddsd_out) + + # Compute dactivation_fn_out. + # + # NOTE: We reuse the activation_fn_out allocation. + dactivation_fn_out = activation_fn_out + stk.backend.triton_kernels.sdd( + ddsd_out, + w2.t(), + dactivation_fn_out.shape, + dactivation_fn_out.data, + dactivation_fn_out.offsets, + dactivation_fn_out.row_indices, + dactivation_fn_out.column_indices, + ) + + # Compute dsdd_out. + # + # NOTE: This reuses the dactivation_fn_out allocation. + if activation_fn is DEFAULT_ACTIVATION_FN: + dsdd_out = gelu.gelu_backward_(dactivation_fn_out, sdd_out) + else: + assert activation_grad_fn is not None + activation_grad_fn(dactivation_fn_out.data) + dsdd_out = stk.Matrix(ctx.shape, sdd_out.data.grad, *topo_tensors) + + # Compute dw1. + dw1 = stk.ops.dsd(dsdd_out.t(), x) + + # Compute dx. + # + # NOTE: This reuses the ddsd_out allocation. + stk.backend.triton_kernels.dsd( + dsdd_out.shape, + dsdd_out.data, + dsdd_out.offsets, + dsdd_out.row_indices, + dsdd_out.column_indices, + dsdd_out.offsets_t, + dsdd_out.column_indices_t, + dsdd_out.block_offsets_t, + False, + w1, + ddsd_out, + ) + dx = ddsd_out + return dx, dw1, dw2, None, None + + +memory_optimized_mlp = MemoryOptimizedMLP.apply + + +class SparseMLP(torch.nn.Module): + + def __init__(self, args: Arguments): + super().__init__() + self.args = args + self._num_rows_per_rank = mpu.experts_per_rank(args) * mpu.features_per_rank(args) + + self.w1 = torch.nn.Parameter( + torch.empty( + self._num_rows_per_rank, + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + self.w2 = torch.nn.Parameter( + torch.empty( + self._num_rows_per_rank, + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + + # Initialize the parameters for the MLP. + # + # NOTE: It is important that we create the weight tensors prior + # to creating the master weights and slicing our the piece for + # this rank. If the master weights are created first the PyTorch + # caching allocator appears to use the same memory block for these + # and the slice which causes large increases in our peak memory + # usage. + with torch.no_grad(): + self.w1.copy_( + create_dmoe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.init_method, + ), + ) + self.w2.copy_( + create_dmoe_expert_weights( + args, + args.moe_num_experts, + args.ffn_hidden_size, + args.hidden_size, + args.output_layer_init_method, + ), + ) + + self._should_set_parallelism_attribute = args.moe_expert_model_parallelism + mpu.set_expert_model_parallel_attributes( + self.w1, + self._should_set_parallelism_attribute, + ) + mpu.set_expert_model_parallel_attributes( + self.w2, + self._should_set_parallelism_attribute, + ) + + self.gradient_scale = None + if self.args.moe_expert_model_parallelism: + self.gradient_scale = 1 / mpu.get_expert_parallel_world_size(self.args,) + + def scale_grad(self, w): + if self.gradient_scale is None: + return w + return scale_gradient(w, self.gradient_scale) + + def forward(self, x, topo): + w1, w2 = self.scale_grad(self.w1), self.scale_grad(self.w2) + w1, w2 = resolve_dtensor(w1), resolve_dtensor(w2) + if self.args.memory_optimized_mlp: + return memory_optimized_mlp( + x, + w1, + w2, + topo, + self.args.activation_fn, + ) + + # Compute the MLP. + x = stk.ops.sdd(x, w1.t(), topo) + activation_fn_out = act_fn(x, self.args.activation_fn) + return stk.ops.dsd(activation_fn_out, w2) + + +class MemoryOptimizedGroupedMLP(torch.autograd.Function): + """GroupedMLP with manually scheduled memory reuse.""" + + @staticmethod + @torch.amp.autocast_mode.custom_fwd(device_type='cuda') + def forward(ctx, x, w1, w2, batch_sizes, activation_fn): + # Cast inputs using ctx dtype from AMP + if ctx._fwd_used_autocast: + x = x.to(ctx._dtype) + w1 = w1.to(ctx._dtype) + w2 = w2.to(ctx._dtype) + # x: [m, k], w1: [n, k], w2: [n, k] + if (not x.is_contiguous() or not w1.is_contiguous() or not w2.is_contiguous()): + raise ValueError("Expected contiguous 'x', 'w1' and 'w2'.") + + # Layer 0: x @ w1.t(). + assert gg.backend is not None + sdd_out = gg.backend.gmm(x, w1, batch_sizes, trans_b=True) + + # activation_fn + activation_fn_out = activation_fn(sdd_out) + + # Layer 1: x @ w2. + dsd_out = gg.backend.gmm(activation_fn_out, w2, batch_sizes) + + # NOTE: Save the input to the layer and the activation_fn input for + # gradient computation. We'll re-compute the activation_fn forward + # pass in the backward pass to avoid materializing another + # intermediate. + ctx.x_shape = x.shape + ctx.sdd_out_shape = sdd_out.shape + ctx.dtype = x.dtype + ctx.activation_fn = activation_fn + ctx.save_for_backward(w1, w2, batch_sizes, x, sdd_out) + return dsd_out + + @staticmethod + @torch.amp.autocast_mode.custom_bwd(device_type='cuda') + def backward(ctx: Any, ddsd_out: torch.Tensor): + if (not ctx.needs_input_grad[0] or not ctx.needs_input_grad[1] or not ctx.needs_input_grad[2]): + raise ValueError('Expected all MLP inputs to need grad.') + + # Unpack saved tensors + # dtype = ctx.dtype + saved_tensors = ctx.saved_tensors + w1, w2 = saved_tensors[:2] + batch_sizes = saved_tensors[2] + x = saved_tensors[3] + sdd_out = saved_tensors[4] + + # Rematerialize activation_fn output. + activation_fn = ctx.activation_fn + with torch.set_grad_enabled(True): + sdd_out.requires_grad = True + activation_fn_out = activation_fn(sdd_out) + activation_grad_fn = activation_fn_out.backward + + # Compute dw2 with recomputed activation_fn output. + assert gg.backend is not None + dw2 = gg.backend.gmm( + activation_fn_out, + ddsd_out, + batch_sizes, + trans_a=True, + ) + + # Compute dactivation_fn_out. + # + # NOTE: We reuse the activation_fn_out allocation. + dactivation_fn_out = activation_fn_out + gg.backend.gmm( + ddsd_out, + w2, + batch_sizes, + trans_b=True, + c=dactivation_fn_out, + ) + + # Compute dsdd_out. + # + # NOTE: This reuses the dactivation_fn_out allocation. + if activation_fn is DEFAULT_ACTIVATION_FN: + dsdd_out = gelu.gelu_backward_(dactivation_fn_out, sdd_out) + else: + assert activation_grad_fn is not None + activation_grad_fn(dactivation_fn_out) + dsdd_out = sdd_out.grad + + # Compute dw1. + dw1 = gg.backend.gmm(dsdd_out, x, batch_sizes, trans_a=True) + + # Compute dx. + # + # NOTE: This reuses the ddsd_out allocation. + gg.backend.gmm(dsdd_out, w1, batch_sizes, c=ddsd_out) + dx = ddsd_out + return dx, dw1, dw2, None, None + + +memory_optimized_grouped_mlp = MemoryOptimizedGroupedMLP.apply + + +class GroupedMLP(SparseMLP): + + def forward(self, x, tokens_per_expert): + batch_sizes = tokens_per_expert.cpu().to(torch.long) + w1, w2 = (self.scale_grad(self.w1), self.scale_grad(self.w2)) + + # Re-shape the weights for the grouped GEMMs. + ne = mpu.experts_per_rank(self.args) + w1 = resolve_dtensor(w1).view(ne, -1, self.args.hidden_size) + w2 = resolve_dtensor(w2).view(ne, -1, self.args.hidden_size) + + if self.args.memory_optimized_mlp: + return memory_optimized_grouped_mlp( + x, + w1, + w2, + batch_sizes, + self.args.activation_fn, + ) + + # Compute the MLP. + assert gg.ops is not None + x = gg.ops.gmm(x, w1, batch_sizes, trans_b=True) + x = self.args.activation_fn(x) + return gg.ops.gmm(x, w2, batch_sizes) + + +class SharedMLP(torch.nn.Module): + """MLP for shared expert. + + Note: this is a copy -> pasta -> modify of the LLM-Foundry MPTMLP class + """ + + def __init__(self, args: Arguments): + super().__init__() + self.args = args + self.fc_kwargs: dict[str, Any] = { + 'bias': args.bias, + 'device': args.device, + } + self.fc_kwargs.update(args.fc_kwargs) + + self.up_proj = args.fc_cls( + args.hidden_size, + args.shared_expert_hidden_size, + **self.fc_kwargs, + ) + self.act = args.activation_fn + self.down_proj = args.fc_cls( + args.shared_expert_hidden_size, + args.hidden_size, + **self.fc_kwargs, + ) + self.down_proj._is_residual = True # a flag for llm-foundry init + + def add_experts_sharedexpert( + self, + shared_expert_out: torch.Tensor, + expert_out: torch.Tensor, + ) -> torch.Tensor: + # Helper function to add expert output to shared expert output + # with optional weighted sum. + if self.args.shared_expert_weighted_sum: + # enable using weighted sum for shared expert output + # wieghted by number of experts used + t_experts = self.args.moe_top_k + 1 + sh_mlp_out = shared_expert_out / t_experts + return sh_mlp_out.add( + expert_out, + alpha=(self.args.moe_top_k / t_experts), + ) + + return shared_expert_out + expert_out + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act(self.up_proj(x))) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/moe.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/moe.py new file mode 100644 index 0000000000000000000000000000000000000000..d0a4aeaacc9c86fc70944e730c53f7a55644e05e --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/moe.py @@ -0,0 +1,507 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Optional, Tuple + +import numpy as np +import torch +import torch.distributed as dist + +# import megablocks.ops as ops +# from megablocks.layers import common, mlp, mpu, router, sharedexpert_registry +# from megablocks.layers.all_to_all import all_to_all +# from megablocks.layers.arguments import Arguments + +from ..ops import ( + sort, + histogram, + inclusive_cumsum, + exclusive_cumsum, + binned_gather, + binned_scatter, + gather, + scatter, + repeat, + replicate, +) + +from . import common, mlp, mpu, router, sharedexpert_registry +from .arguments import Arguments +from .all_to_all import all_to_all + +_LOAD_BALANCING_LOSS = [] + + +def save_load_balancing_loss(loss): + global _LOAD_BALANCING_LOSS + _LOAD_BALANCING_LOSS.append(loss) + + +def get_load_balancing_loss(): + global _LOAD_BALANCING_LOSS + return _LOAD_BALANCING_LOSS + + +def clear_load_balancing_loss(): + global _LOAD_BALANCING_LOSS + _LOAD_BALANCING_LOSS.clear() + + +def batched_load_balancing_loss(args: Arguments): + if args.moe_loss_weight == 0: + return 0.0 + + # tokens_per_expert[i].shape = (num_experts) + # expert_scores[i].shape = (tokens, num_experts) + tokens_per_expert, expert_scores = zip(*get_load_balancing_loss()) + num_layers_per_pipeline_stage = (args.num_layers // args.pipeline_model_parallel_size) + if args.num_layers_per_virtual_pipeline_stage is not None: + num_layers_per_pipeline_stage = args.num_layers_per_virtual_pipeline_stage + + if len(tokens_per_expert) != num_layers_per_pipeline_stage: + raise ValueError( + f'Expected {num_layers_per_pipeline_stage} token_per_experts ' + f'but found {len(tokens_per_expert)}.\nnum_layers = ' + f'{args.num_layers}\npipeline_model_parallel_size = ' + f'{args.pipeline_model_parallel_size}\n' + 'num_layers_per_virtual_pipeline_stage' + f' = {args.num_layers_per_virtual_pipeline_stage}', + ) + if len(expert_scores) != num_layers_per_pipeline_stage: + raise ValueError( + f'Expected {num_layers_per_pipeline_stage} expert_scores ' + f'but found {len(tokens_per_expert)}.\nnum_layers = ' + f'{args.num_layers}\npipeline_model_parallel_size = ' + f'{args.pipeline_model_parallel_size}\n' + 'num_layers_per_virtual_pipeline_stage' + f' = {args.num_layers_per_virtual_pipeline_stage}', + ) + + # Verify the shape of the tokens_per_expert and expert_scores tensors. + assert all((x.ndim == 1 and x.numel() == args.moe_num_experts for x in tokens_per_expert)) + + tokens = expert_scores[0].shape[0] + assert all(((x.ndim == 2 and x.shape[1] == args.moe_num_experts and x.shape[0] == tokens) for x in expert_scores)) + + # Concatenate the contributions of each layer and convert to + # the correct types and formats for the dot product. + expert_scores = torch.cat(expert_scores, dim=1) + if args.moe_lbl_in_fp32: + expert_scores = expert_scores.float() + if tokens != 0: + expert_scores = expert_scores.mean(dim=0) + else: + expert_scores = expert_scores.sum(dim=0) + tokens_per_expert = torch.cat(tokens_per_expert).to(expert_scores.dtype) + + expected_values = num_layers_per_pipeline_stage * args.moe_num_experts + assert tokens_per_expert.numel() == expected_values + assert expert_scores.numel() == expected_values + + # Calculate the total scale across all factors. + # + # loss_weight * num_experts / (num_layers * tokens * top_k) + scale_numerator = (args.moe_num_experts * args.moe_loss_weight) + scale_denominator = (args.num_layers * tokens * args.moe_top_k) + scale = scale_numerator / scale_denominator + return scale * torch.dot(tokens_per_expert, expert_scores) + + +# NOTE: This class defines MoE expert computation, including expert model parallel +# communication. When using FSDP on top of MegaBlocks this is the module that should +# be wrapped s.t. the weight all-gathers can be scheduled *before* the expert model +# parallel all2all. +class ParallelMLP(torch.nn.Module): + + def __init__(self, args: Arguments): + super(ParallelMLP, self).__init__() + self.args = args + + # Calculate the number of experts in total and the number of experts + # owned by this rank. + # world_size = mpu.get_expert_parallel_world_size(args) + self.num_experts = args.moe_num_experts + self.top_k = self.args.moe_top_k + + # Calculate the number of bits needed to represent the expert indices + # so that we can pass it to radix sort. + self.sort_end_bit = max(int(np.ceil(np.log2(self.num_experts))), 1) + + # Expert MLP. + self.mlp = mlp.MLP(args) + + self.bias: Optional[torch.Tensor] + if self.args.bias: + # Note that the output bias is not parallelized with expert + # model parallelism. + self.bias = torch.nn.Parameter( + torch.empty( + args.hidden_size, + device=args.device, + dtype=common.dtype(args), + ), + ) + torch.nn.init.zeros_(self.bias) + else: + self.register_parameter('bias', None) + + # Select the forward function for the operating mode. + self.forward_fn = (self.parallel_forward_once if args.moe_expert_model_parallelism else self.forward_once) + + def expert_capacity(self, tokens: int) -> int: + world_size = mpu.get_expert_parallel_world_size(self.args) + tokens_per_expert = (self.top_k * tokens * world_size / self.num_experts) + return int(self.args.moe_capacity_factor * tokens_per_expert) + + def load_balancing_loss(self, tokens_per_expert: torch.Tensor, expert_scores: torch.Tensor): + """Calculate the load balancing loss contribution.""" + assert len(expert_scores.size()) == 2 + tokens, num_experts = expert_scores.size() + assert num_experts == self.num_experts + assert len(tokens_per_expert.size()) == 1 + num_experts, = tokens_per_expert.size() + assert num_experts == self.num_experts + scale = self.num_experts / (tokens * self.top_k) + return scale * torch.dot( + tokens_per_expert.to(expert_scores.dtype), + expert_scores.mean(dim=0), + ) + + def indices_and_bins(self, + top_expert: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # Sort the expert ids to produce the scatter/gather + # indices for the permutation. + # + # TODO(tgale): Is it worth doing this conversion to 32-bit + # prior? Could we place the `torch.max` operation to return + # 32-bit expert indices? + top_expert = top_expert.int() + # output = ops.sort(top_expert, self.sort_end_bit) + output = sort(top_expert, self.sort_end_bit) + assert output is not None + bin_ids, indices = output + + # Histogram the expert ids to identify the number of + # tokens routed to each expert. + # + # TODO(tgale): Does the sorted data produce a more favorable + # data distribution for histogram? Or is the op parallelism + # worth more? + # tokens_per_expert = ops.histogram(top_expert, self.num_experts) + tokens_per_expert = histogram(top_expert, self.num_experts) + + # Calculate the bin bounds for the sorted tokens. + # bins = ops.inclusive_cumsum(tokens_per_expert, 0) + bins = inclusive_cumsum(tokens_per_expert, 0) + assert bins is not None + bins = bins.view(1) if not len(bins.size()) else bins + + assert isinstance(indices, torch.Tensor) + assert isinstance(bin_ids, torch.Tensor) + assert isinstance(bins, torch.Tensor) + assert isinstance(tokens_per_expert, torch.Tensor) + + return indices, bin_ids, bins, tokens_per_expert + + def permute_and_compute( + self, + x: torch.Tensor, + tokens_per_expert: int, # unused + indices: torch.Tensor, + bin_ids: torch.Tensor, # unused + expert_weights: torch.Tensor, + bins: torch.Tensor, + expert_capacity: int, + top_k: int, + ): + # Route the tokens for MoE computation. + x = x.view(-1, x.shape[-1]) + # output = ops.binned_gather(x, indices, bins, expert_capacity, top_k) + output = binned_gather(x, indices, bins, expert_capacity, top_k) + assert output is not None + x = output + + # Perform the expert computation. Note that we don't + # use biases for these linear operations. + x = self.mlp(x) + + # Un-route the data for the MoE output. + # return ops.binned_scatter(x, indices, expert_weights, bins, top_k) + return binned_scatter(x, indices, expert_weights, bins, top_k) + + + def forward_once(self, x: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor): + # x: [sl, bs, hs] + # expert_weights: [sl * bs, top-k] + # top_experts: [sl * bs, top-k] + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + with torch.no_grad(): + indices, bin_ids, bins, tokens_per_expert = (self.indices_and_bins(top_experts)) + + # If expert_capacity is set to zero, set the number of tokens + # per expert to the maximum we need to avoid dropping tokens. + sl, bs, _ = x.size() + expert_capacity = self.expert_capacity(sl * bs) + if expert_capacity == 0: + expert_capacity = torch.max(tokens_per_expert).item() + + x = self.permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capacity, + self.top_k, + ) + return x, tokens_per_expert + + def parallel_forward_once(self, x: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor): + # NOTE: This function implements the same computation as forward_once + # but with expert model parallelism. + # + # 1. Permute the tokens locally so that they are grouped by their + # expert assignments. This allows us to transfer all of the tokens + # for a remote device in one communication primitive. + # + # 2. Permute the tokens across the expert parallel devices. After + # this is completed each device has all of the tokens assigned to + # its set of experts in its local HBM. + # + # 3. Permute the tokens locally so that they are grouped by their + # expert assignement. After the distributed permutation the tokens + # are grouped by which device they came from. We re-order them + # locally to allow for efficient computation. + # + # After this series of permutations we compute the linear layers + # and then repeat these three steps in reverse to produce the final + # output. + # + # Compute the mapping of local tokens to experts. + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + with torch.no_grad(): + indices, bin_ids, bins, tokens_per_expert = (self.indices_and_bins(top_experts)) + + # If we're sharding the experts along the hidden dimension + # multiple devices own parts of the same sets of experts. + # Replicate the token counts so every device gets the counts. + # repeated_tokens_per_expert = ops.repeat( + repeated_tokens_per_expert = repeat( + tokens_per_expert, + (mpu.hidden_sharding_degree(self.args),), + ) + + # Pass token count information to the device on which the + # target expert resides. + parallel_tokens_per_expert = torch.empty_like(repeated_tokens_per_expert,) + tpe_handle = dist.all_to_all_single( + parallel_tokens_per_expert, + repeated_tokens_per_expert, + group=self.args.expert_parallel_group, + async_op=True, + ) + + # Permute locally and without any padding so that tokens for each + # parallel device are stored contiguously. + # + # This view updates the shape of the tensor from [sl, bs, hs] to + # [sl * bs, hs] prior to the permutation. + x = x.view(-1, x.shape[-1]) + # output = ops.gather(x, indices, bin_ids, bins, self.top_k) + output = gather(x, indices, bin_ids, bins, self.top_k) + assert output is not None + x = output + + # Compute the number of tokens that will be received from each + # device and permute the input data across the devices. + with torch.no_grad(): + tpe_handle.wait() + experts_per_rank = mpu.experts_per_rank(self.args) + + # Reshape to [world_size, num_experts_per_rank]. + world_size = mpu.get_expert_parallel_world_size(self.args) + repeated_tokens_per_expert = (repeated_tokens_per_expert.view(world_size, experts_per_rank)) + parallel_tokens_per_expert = (parallel_tokens_per_expert.view(world_size, experts_per_rank)) + + # TODO(tgale): It might be faster to do this on the GPU and + # then communicate the results back to the host. + send_counts = repeated_tokens_per_expert.cpu().sum(dim=-1) + parallel_tokens_per_expert_cpu = parallel_tokens_per_expert.cpu() + recv_counts = parallel_tokens_per_expert_cpu.sum(dim=-1) + + # Convert the send/recv counts to lists. + send_counts = send_counts.tolist() + recv_counts = recv_counts.tolist() + tokens_received = sum(recv_counts) + + # If we're sharding the experts along the hidden dimension + # multiple devices own parts of the same sets of experts. + # Replicate the token counts so devices that share experts + # get all of the tokens assigned to them. + # + # TODO(tgale): Fuse this into the prior, local permutation. + # x = ops.repeat(x, (mpu.hidden_sharding_degree(self.args), 1)) + x = repeat(x, (mpu.hidden_sharding_degree(self.args), 1)) + + # Start the cross-device permutation asynchronously so we can + # overlap communication with computation. + parallel_x, parallel_x_handle = all_to_all( + x, + recv_counts, + send_counts, + self.args.expert_parallel_group, + async_op=True, + ) + + with torch.no_grad(): + # After we do the cross-device permutation we have the tokens on the + # correct device but not yet grouped by expert because we received + # tokens from each device as contiguous chunks. To group the tokens + # for expert computation we'll do one more local permutation. The + # rest of this torch.no_grad() scope sets up the indices and bins + # for this permutation. + # replicate_bins = ops.inclusive_cumsum( + replicate_bins = inclusive_cumsum( + parallel_tokens_per_expert.flatten(), + 0, + ) + replicate_bins = (replicate_bins.view(1) if not len(replicate_bins.size()) else replicate_bins) + + # Construct the expert indices for the permuted tokens. + parallel_top_expert = torch.remainder( + torch.arange( + self.num_experts * mpu.hidden_sharding_degree(self.args), + dtype=torch.int32, + device=indices.device, + ), + mpu.experts_per_rank(self.args), + ) + # parallel_top_expert = ops.replicate( + parallel_top_expert = replicate( + parallel_top_expert.unsqueeze(dim=0), + replicate_bins, + tokens_received, + ).flatten() + + # TODO(tgale): The sort_end_bit here can be reduced. + # parallel_bin_ids, parallel_indices = ops.sort( + parallel_bin_ids, parallel_indices = sort( + parallel_top_expert, + self.sort_end_bit, + ) + + # Calculate the bins boundaries from the token counts. + parallel_tokens_per_expert = parallel_tokens_per_expert.sum( + dim=0, + dtype=torch.int, + ) + # parallel_bins = ops.inclusive_cumsum(parallel_tokens_per_expert, 0) + parallel_bins = inclusive_cumsum(parallel_tokens_per_expert, 0) + parallel_bins = (parallel_bins.view(1) if not len(parallel_bins.size()) else parallel_bins) + + # If expert_capacity is set to zero, set the number of tokens + # per expert to the maximum we need to avoid dropping tokens. + tokens, _ = x.size() + expert_capacity = self.expert_capacity(tokens) + if expert_capacity == 0: + expert_capacity = torch.max(parallel_tokens_per_expert).item() + + # Locally permute the tokens and perform the expert computation. + # Block to make sure that the cross-device permutation is complete. + if self.args.mlp_impl == 'grouped': + # GroupedMLP requires counts on CPU. We can use the tensor already + # moved to CPU for the prior all_to_all, which avoids an extra + # device synchronization. + parallel_tokens_per_expert = parallel_tokens_per_expert_cpu.sum( + dim=0, + dtype=torch.int, + ) + parallel_x_handle.wait() + parallel_x = self.permute_and_compute( + parallel_x, + parallel_tokens_per_expert, + parallel_indices, + parallel_bin_ids, + None, # expert_weights + parallel_bins, + expert_capacity, + top_k=1, + ) + + # Un-permute the tokens across the devices. + x, _ = all_to_all( + parallel_x, + send_counts, + recv_counts, + self.args.expert_parallel_group, + ) + + # Reduce along the hidden sharding to get the final outputs. + # + # TODO(tgale): Fuse this into the following local permutation. + shape = ( + mpu.hidden_sharding_degree(self.args), + -1, + self.args.hidden_size, + ) + # x = ops.sum(x.view(shape), dim=0) + x = x.view(shape).sum(dim=0) + + # Un-permute locally to setup for the next series of operations. + # x = ops.scatter(x, indices, bin_ids, expert_weights, bins, self.top_k) + x = scatter(x, indices, bin_ids, expert_weights, bins, self.top_k) + return x, tokens_per_expert.flatten() + + def forward(self, x: torch.Tensor, scores: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor): + in_shape = x.size() + + # Compute the experts. + x, tokens_per_expert = self.forward_fn(x, expert_weights, top_experts) + if self.training and self.args.moe_loss_weight > 0: + save_load_balancing_loss((tokens_per_expert, scores)) + x = x.view(in_shape) + if self.bias is not None: + if self.args.return_bias: + return x, self.bias + return x + self.bias + return x + + +class MoE(torch.nn.Module): + + def __init__(self, args: Arguments): + super(MoE, self).__init__() + + # Token router. + self.router = router.LearnedRouter(args) + + # Expert computation helper. + self.experts = self._init_experts_mlp(args) + + self.shared_expert = None + if args.shared_expert: + # SharedExpert computation helper. + self.shared_expert = sharedexpert_registry.get(args) + + def _init_experts_mlp(self, args: Arguments): + return ParallelMLP(args) + + def forward(self, x: torch.Tensor): + # NOTE: If we're going to cast the activations to lower precision + # do it before we permute the tokens to save bandwidth. + x = common.cast_if_autocast_enabled(x) + + # Compute the expert scores and assignments. + scores, expert_weights, top_experts = self.router(x) + + # Compute the experts. + out = self.experts(x, scores, expert_weights, top_experts) + if self.shared_expert is not None: + shared_expert_out = self.shared_expert(x) + out = self.shared_expert.add_experts_sharedexpert( + shared_expert_out, + out, + ) + return out diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/mpu.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/mpu.py new file mode 100644 index 0000000000000000000000000000000000000000..434e143ab42bf3f83406d69e9dd1f72777716e22 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/mpu.py @@ -0,0 +1,94 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Optional + +import torch +import torch.distributed as dist + +# from megablocks.layers.arguments import Arguments +from .arguments import Arguments + + +class MoeParam(torch.Tensor): + + def __init__(self): + super().__init__(self) + self.expert_model_parallel: bool + + +def is_moe_param(tensor: torch.Tensor) -> bool: + return hasattr(tensor, 'expert_model_parallel') + + +def get_expert_parallel_world_size(args: Arguments) -> int: + return (dist.get_world_size(args.expert_parallel_group) if args.moe_expert_model_parallelism else 1) + + +def get_expert_parallel_rank(args: Arguments) -> int: + return (dist.get_rank(args.expert_parallel_group) if args.moe_expert_model_parallelism else 0) + + +def set_expert_model_parallel_attributes( + tensor: torch.Tensor, + is_parallel: bool, +): + assert not hasattr(tensor, 'expert_model_parallel') + setattr(tensor, 'expert_model_parallel', is_parallel) + + +def param_is_expert_model_parallel(param: MoeParam) -> bool: + return (hasattr(param, 'expert_model_parallel') and param.expert_model_parallel) + + +def copy_expert_model_parallel_attributes( + destination_tensor: torch.Tensor, + source_tensor: torch.Tensor, +): + if hasattr(source_tensor, 'expert_model_parallel'): + setattr( + destination_tensor, + 'expert_model_parallel', + getattr(source_tensor, 'expert_model_parallel'), + ) + + +def synchronized_print(group: Optional[dist.ProcessGroup], *x: torch.Tensor): + world_size = dist.get_world_size(group) + rank = dist.get_rank(group) + for i in range(world_size): + dist.barrier(group) + if i == rank: + print(f'rank = {rank}', *x) + + +# Helpers for expert/tensor sharding. +def expert_sharding_degree(args: Arguments) -> int: + world_size = get_expert_parallel_world_size(args) + esd = min(world_size, args.moe_num_experts) + + if (args.moe_num_experts % esd) != 0: + raise ValueError(f'Cannot shard {args.moe_num_experts} experts {esd} ways.',) + return esd + + +def hidden_sharding_degree(args: Arguments) -> int: + world_size = get_expert_parallel_world_size(args) + esd = expert_sharding_degree(args) + hsd = world_size // esd + + if (args.ffn_hidden_size % hsd) != 0: + raise ValueError(f'Cannot shard {args.ffn_hidden_size} features {hsd} ways.',) + if (esd * hsd) != world_size: + raise ValueError( + f"Invalid sharding. 'expert_sharding_degree' ({esd}) * hidden_sharding_degree ({hsd}) != world_size ({world_size}).", + ) + return hsd + + +def experts_per_rank(args: Arguments) -> int: + return args.moe_num_experts // expert_sharding_degree(args) + + +def features_per_rank(args: Arguments) -> int: + return args.ffn_hidden_size // hidden_sharding_degree(args) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/router.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/router.py new file mode 100644 index 0000000000000000000000000000000000000000..37cb2782348d62583376f1a183c7ede83601216d --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/router.py @@ -0,0 +1,116 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch + +# from megablocks.layers import common +# from megablocks.layers.arguments import Arguments +from . import common +from .arguments import Arguments + +_ROUTER_LOGITS = [] + + +def _save_router_logits(logits: torch.Tensor, args: Arguments): + if args.moe_zloss_weight == 0: + return + global _ROUTER_LOGITS + _ROUTER_LOGITS.append(logits) + + +def clear_router_zloss(): + global _ROUTER_LOGITS + _ROUTER_LOGITS.clear() + + +def batched_router_zloss(args: Arguments): + global _ROUTER_LOGITS + + if args.moe_zloss_weight == 0: + import warnings + warnings.warn('Call to batched_router_zloss, but moe_zloss_weight=0') + return 0 + + logits_per_router = _ROUTER_LOGITS + + if args.moe_zloss_in_fp32: + logits_per_router = [logits.float() for logits in logits_per_router] + + unscaled_zloss_per_router = torch.stack([ + torch.logsumexp(logits, dim=1).square().mean() for logits in logits_per_router + ]) + + return args.moe_zloss_weight * unscaled_zloss_per_router + + +# NOTE: To enable end-to-end benchmarking without convergence we +# support a flag to force the router to assign tokens uniformly +# across the experts. We do this with a custom autograd operation +# so that PyTorch still executes the full set of router operation. +class _UniformExpertAssignment(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, num_experts: int): + out = torch.arange(x.numel(), dtype=x.dtype, device=x.device) + out = torch.remainder(out, num_experts) + return out.view(x.shape) + + +_uniform_expert_assignment = _UniformExpertAssignment.apply + + +class LearnedRouter(torch.nn.Module): + + def __init__(self, args: Arguments): + super().__init__() + self.args = args + + # Learned router parameters. + # + # NOTE: This weight matrix is not parallelized with expert model + # parallelism. Each device needs the entire router weight matrix + # so that it can route its batch of data correctly. + self.layer = torch.nn.Linear( + args.hidden_size, + args.moe_num_experts, + bias=False, + dtype=common.dtype(args), + device=args.device, + ) + args.init_method(self.layer.weight) + + def jitter(self, x: torch.Tensor): + low: float = 1.0 - self.args.moe_jitter_eps + high: float = 1.0 + self.args.moe_jitter_eps + noise = torch.rand(x.size(), dtype=x.dtype, device=x.device) + return low + noise * (high - low) + + def _top_k(self, scores: torch.Tensor): + if self.args.moe_top_k == 1: + return scores.max(dim=-1, keepdim=True) + return torch.topk(scores, self.args.moe_top_k, dim=-1) + + def forward(self, x: torch.Tensor): + if self.training and self.args.moe_jitter_eps is not None: + x = x * self.jitter(x) + + logits = self.layer(x.view(-1, x.shape[-1])) + _save_router_logits(logits, self.args) + scores = logits.softmax(dim=-1) + expert_weights, expert_indices = self._top_k(scores) + if self.args.moe_normalize_expert_weights: + expert_weights = expert_weights / torch.norm( + expert_weights, + p=self.args.moe_normalize_expert_weights, + dim=-1, + keepdim=True, + ) + + expert_indices = ( + _uniform_expert_assignment( + expert_indices, + self.args.moe_num_experts, + ) if self.args.uniform_expert_assignment else expert_indices + ) + return scores, expert_weights, expert_indices diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_layers/sharedexpert_registry.py b/build/torch214-cxx11-cu132-x86_64-linux/_layers/sharedexpert_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..5840862f88f370ace5fd49bd0612fc98d186cc49 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_layers/sharedexpert_registry.py @@ -0,0 +1,32 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Union + +# from megablocks.layers import glu, mlp +# from megablocks.layers.arguments import Arguments +from . import glu, mlp +from .arguments import Arguments + +_REGISTRY = { + 'mlp': mlp.SharedMLP, + 'glu': glu.SharedGLU, +} + + +def get(args: Arguments) -> Union[mlp.SharedMLP, glu.SharedGLU]: + """Returns an SharedMLP for use in a dMoE instance. + + Uses the provided arguments to instantiate the appropriate + SharedMLP instance. + + Args: + args: propagated Arguments dataclass. + + Returns: + An instantiated SharedMLP constructed using the input args. + """ + if args.mlp_type not in _REGISTRY: + raise ValueError(f'Unsupported mlp type: {args.mlp_type}') + + return _REGISTRY[args.mlp_type](args) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_megablocks_cuda_addf474.abi3.so b/build/torch214-cxx11-cu132-x86_64-linux/_megablocks_cuda_addf474.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..d318eb1aa38134a1314696163191bca466e69d0c --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_megablocks_cuda_addf474.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3c87c27c4d7de338dfb1a3eca4f63fb8585ba2f144458c3bd1966e4616e2e462 +size 14987032 diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_ops.py b/build/torch214-cxx11-cu132-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..6640ac69d26eee39b52fe070e3c8571ed6a93d4e --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _megablocks_cuda_addf474 +ops = torch.ops._megablocks_cuda_addf474 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_megablocks_cuda_addf474::{op_name}" diff --git a/build/torch214-cxx11-cu132-x86_64-linux/_version.py b/build/torch214-cxx11-cu132-x86_64-linux/_version.py new file mode 100644 index 0000000000000000000000000000000000000000..c55783177af19bc03654c730c4892df8f8532279 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/_version.py @@ -0,0 +1,6 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +"""The MegaBlocks Version.""" + +__version__ = '0.11.0.dev0' diff --git a/build/torch214-cxx11-cu132-x86_64-linux/backend/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/backend/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9d4e43e9b7e6a9a1c3bde2df34914643ca5d8332 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/backend/__init__.py @@ -0,0 +1,2 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 diff --git a/build/torch214-cxx11-cu132-x86_64-linux/backend/kernels.py b/build/torch214-cxx11-cu132-x86_64-linux/backend/kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..d20139324883338ddb312e4b05a72056d74491ac --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/backend/kernels.py @@ -0,0 +1,557 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch +import triton +import triton.language as tl + +# Stub triton autotune when testing in a env that does not have CUDA +# this approach preserves the original code but enables testing without a GPU +if torch.cuda.is_available() is False: + import warnings + + warnings.warn("CUDA is not available. Triton autotuning is disabled.") + + def _no_autotune(*args, **kwargs): + def deco(fn): + return fn + return deco + + triton.autotune = _no_autotune + + +def assert_is_tensor(x, ndim): + if x.ndim != ndim: + raise ValueError(f'Expected {ndim}-tensor but got {x.ndim}-tensor') + + +def assert_is_matrix(x): + assert_is_tensor(x, 2) + + +def assert_is_vector(x): + if x.ndim != 1: + raise ValueError(f'Expected 1-tensor but got {x.ndim}-tensor') + + +def assert_equal(a, b): + if a != b: + raise ValueError(f'Expected dimensions to be equal but got {a} and {b}.',) + + +# a: (tokens, hidden_size), real. +# indices: (tokens * top_k), integer. +# bin_ids: (tokens * top_k), integer. +# weights: (tokens * top_k), real. +# bins: (num_experts), integer. +# padded_bins: (num_experts), integer. +@triton.autotune( + configs=[ + triton.Config({'BLOCK_X': 64}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=2), + triton.Config({'BLOCK_X': 256}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=4), + triton.Config({'BLOCK_X': 256}, num_warps=4), + ], + key=['NUM_COLUMNS'], +) +@triton.jit +def _padded_copy( + a, + b, + indices, + bin_ids, + weights, + bins, + padded_bins, + NUM_COLUMNS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_X: tl.constexpr, + A_TO_B: tl.constexpr, + SCALE: tl.constexpr, +): + # Our index into array 'a'. + index_a = tl.load(indices + tl.program_id(0)) + + # One threadblock per row in 'a'. Array 'b' has greater or equal + # number of rows since they could be padded. + bin_idx = tl.load(bin_ids + tl.program_id(0)) + + # Now we know what bin we're assigned to, but we need to know how + # many threadblocks were assigned to earlier bins so we can offset + # in our bin properly. + offset_in_bin = tl.program_id(0) + if bin_idx > 0: + offset_in_bin -= tl.load(bins + bin_idx - 1) + + # Load the starting index of our bin in array 'b'. + index_b = offset_in_bin + if bin_idx > 0: + index_b += tl.load(padded_bins + bin_idx - 1) + + # Offset the input and output pointers. + # + # If we're going from A to B, divide the input index to copy + # the same input repeatedly. If we're going from B to A we + # need to reduce the result. Using atomics is slow, so we + # do the reduce step in a second kernel. + offset = index_a // TOP_K if A_TO_B else index_a + a += tl.multiple_of(offset * NUM_COLUMNS, NUM_COLUMNS) + b += tl.multiple_of(index_b * NUM_COLUMNS, NUM_COLUMNS) + offsets = tl.max_contiguous(tl.arange(0, BLOCK_X), BLOCK_X) + + # Load the scale, if requested. + scale = tl.load(weights + index_a) if SCALE else 1 + + # Swap the pointers depending on the direction. + iptr = a if A_TO_B else b + optr = b if A_TO_B else a + + iterations = tl.cdiv(NUM_COLUMNS, BLOCK_X) + for _ in range(iterations): + mask = offsets < NUM_COLUMNS + x = tl.load(iptr + offsets, mask=mask) + x = x.to(tl.float32) * scale.to(tl.float32) + + tl.store(optr + offsets, x.to(optr.dtype.element_ty), mask=mask) + + offsets += BLOCK_X + + +def padded_gather(x, indices, bin_ids, weights, bins, padded_bins, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_vector(indices) + assert_is_vector(bin_ids) + assert_is_vector(bins) + assert_is_vector(padded_bins) + assert_equal(indices.shape[0], x.shape[0] * top_k) + assert_equal(bin_ids.shape[0], x.shape[0] * top_k) + assert_equal(bins.size(), padded_bins.size()) + + if weights is not None: + assert_equal(weights.shape[0], x.shape[0] * top_k) + + # NOTE: Because of the padding, the output size is dynamic. + # We load the final padded bin bound to get the output rows. + output_rows = padded_bins[-1].cpu().item() + out = torch.zeros((output_rows, x.shape[1]), dtype=x.dtype, device=x.device) + _padded_copy[(indices.shape[0],)]( + x, + out, + indices, + bin_ids, + weights, + bins, + padded_bins, + NUM_COLUMNS=x.shape[1], + A_TO_B=True, + TOP_K=top_k, + SCALE=weights is not None, + ) + return out + + +def gather(x, indices, bin_ids, weights, bins, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_vector(indices) + assert_is_vector(bin_ids) + assert_is_vector(bins) + assert_equal(indices.shape[0], x.shape[0] * top_k) + assert_equal(bin_ids.shape[0], x.shape[0] * top_k) + + if weights is not None: + assert_equal(weights.shape[0], x.shape[0] * top_k) + + # NOTE: There is no padding so the output rows equals the + # input rows multiplied by top_k. + output_rows = x.shape[0] * top_k + out = torch.empty((output_rows, x.shape[1]), dtype=x.dtype, device=x.device) + _padded_copy[(indices.shape[0],)]( + x, + out, + indices, + bin_ids, + weights, + bins, + bins, + NUM_COLUMNS=x.shape[1], + A_TO_B=True, + TOP_K=top_k, + SCALE=weights is not None, + ) + return out + + +def padded_scatter(x, indices, bin_ids, weights, bins, padded_bins, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_vector(indices) + assert_is_vector(bin_ids) + assert_is_vector(bins) + assert_is_vector(padded_bins) + assert_equal(indices.shape[0], bin_ids.shape[0]) + assert_equal(bins.size(), padded_bins.size()) + + if weights is not None: + assert_equal(indices.shape[0], weights.shape[0]) + + tokens = indices.shape[0] // top_k + out = torch.empty((tokens, top_k, x.shape[1]), dtype=x.dtype, device=x.device) + _padded_copy[(indices.shape[0],)]( + out, + x, + indices, + bin_ids, + weights, + bins, + padded_bins, + NUM_COLUMNS=x.shape[1], + A_TO_B=False, + TOP_K=top_k, + SCALE=weights is not None, + ) + + # Reduce along the top-k dimension, if needed. + return out.sum(dim=1) if top_k > 1 else out.view(tokens, x.shape[1]) + + +def scatter(x, indices, bin_ids, weights, bins, top_k): + return padded_scatter(x, indices, bin_ids, weights, bins, bins, top_k) + + +# x: (tokens, top_k, hidden_size), real +# grad: (tokens, hidden_size), real. +# wgrad: (tokens, top_k), real. +# indices: (tokens * top_k), integer. +# bin_ids: (tokens * top_k), integer. +# bins: (num_experts), integer. +# padded_bins: (num_experts), integer. +@triton.autotune( + configs=[ + triton.Config({'BLOCK_X': 64}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=2), + triton.Config({'BLOCK_X': 256}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=4), + triton.Config({'BLOCK_X': 256}, num_warps=4), + ], + key=['NUM_COLUMNS'], +) +@triton.jit +def _padded_copy_wgrad( + x, + grad, + wgrad, + indices, + bin_ids, + bins, + padded_bins, + NUM_COLUMNS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_X: tl.constexpr, +): + # Our index into 'tokens * top_k'. + index_out = tl.load(indices + tl.program_id(0)) + + # One threadblock per row in 'a'. Array 'b' has greater or equal + # number of rows since they could be padded. + bin_idx = tl.load(bin_ids + tl.program_id(0)) + + # Now we know what bin we're assigned to, but we need to know how + # many threadblocks were assigned to earlier bins so we can offset + # in our bin properly. + offset_in_bin = tl.program_id(0) + if bin_idx > 0: + offset_in_bin -= tl.load(bins + bin_idx - 1) + + # Load the starting index of our bin in array 'x'. + index_x = offset_in_bin + if bin_idx > 0: + index_x += tl.load(padded_bins + bin_idx - 1) + + # Offset the input and output pointers. + wgrad += index_out + grad += tl.multiple_of((index_out // TOP_K) * NUM_COLUMNS, NUM_COLUMNS) + x += tl.multiple_of(index_x * NUM_COLUMNS, NUM_COLUMNS) + offsets = tl.max_contiguous(tl.arange(0, BLOCK_X), BLOCK_X) + + acc = tl.zeros((BLOCK_X,), dtype=tl.float32) + iterations = tl.cdiv(NUM_COLUMNS, BLOCK_X) + for _ in range(iterations): + mask = offsets < NUM_COLUMNS + data = tl.load(x + offsets, mask=mask).to(tl.float32) + scale = tl.load(grad + offsets, mask=mask).to(tl.float32) + acc += data * scale + offsets += BLOCK_X + + # Reduce to get the final result and store. + out = tl.sum(acc).to(wgrad.dtype.element_ty) + tl.store(wgrad, out) + + +def padded_scatter_wgrad(x, grad, indices, bin_ids, bins, padded_bins, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_matrix(grad) + assert_is_vector(indices) + assert_is_vector(bin_ids) + assert_is_vector(bins) + assert_is_vector(padded_bins) + assert_equal(indices.shape[0], bin_ids.shape[0]) + assert_equal(bins.size(), padded_bins.size()) + + tokens = indices.shape[0] // top_k + out = torch.empty((tokens * top_k), dtype=x.dtype, device=x.device) + _padded_copy_wgrad[(indices.shape[0],)]( + x, + grad, + out, + indices, + bin_ids, + bins, + padded_bins, + NUM_COLUMNS=x.shape[1], + TOP_K=top_k, + ) + return out + + +def scatter_wgrad(x, grad, indices, bin_ids, bins, top_k): + return padded_scatter_wgrad(x, grad, indices, bin_ids, bins, bins, top_k) + + +# a: (tokens, hidden_size), real. +# b: (num_experts, expert_capacity, num_columns), real. +# indices: (tokens * top_k), integer. +# weights: (tokens * top_k), real. +# bins: (num_experts), integer. +@triton.autotune( + configs=[ + triton.Config({'BLOCK_X': 64}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=2), + triton.Config({'BLOCK_X': 256}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=4), + triton.Config({'BLOCK_X': 256}, num_warps=4), + ], + key=['NUM_COLUMNS'], +) +@triton.jit +def _binned_copy( + a, + b, + num_experts, + expert_capacity, + indices, + weights, + bins, + NUM_COLUMNS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_X: tl.constexpr, + A_TO_B: tl.constexpr, + SCALE: tl.constexpr, +): + # Load our indices into the output. + expert_idx = tl.program_id(0) + entry_idx = tl.program_id(1) + + # Calculate our offset into the output. + index_b = expert_idx * expert_capacity + entry_idx + + # Load the index bounds for our bin and calculate + # the number of tokens assigned to our expert. + start = 0 + if expert_idx > 0: + start = tl.load(bins + expert_idx - 1) + end = tl.load(bins + expert_idx) + num_tokens = end - start + + # Calculate our offset into the input. If we don't + # have an input exit early. + if entry_idx >= num_tokens: + return + index_a = tl.load(indices + start + entry_idx) + + # Offset the input and output pointers. + # + # If we're going from A to B, divide the input index to copy + # the same input repeatedly. If we're going from B to A we + # need to reduce the result. Using atomics is slow, so we + # do the reduce step in a second kernel. + offset = index_a // TOP_K if A_TO_B else index_a + a += tl.multiple_of(offset * NUM_COLUMNS, NUM_COLUMNS) + b += tl.multiple_of(index_b * NUM_COLUMNS, NUM_COLUMNS) + offsets = tl.max_contiguous(tl.arange(0, BLOCK_X), BLOCK_X) + + # Load the scale, if requested. + scale = tl.load(weights + index_a) if SCALE else 1 + + # Swap the pointers depending on the direction. + # + # NOTE: We need to zero the output in both directions. + iptr = a if A_TO_B else b + optr = b if A_TO_B else a + + iterations = tl.cdiv(NUM_COLUMNS, BLOCK_X) + for _ in range(iterations): + mask = offsets < NUM_COLUMNS + x = tl.load(iptr + offsets, mask=mask) + x = x.to(tl.float32) * scale.to(tl.float32) + + tl.store(optr + offsets, x.to(optr.dtype.element_ty), mask=mask) + + offsets += BLOCK_X + + +def binned_gather(x, indices, weights, bins, expert_capacity, top_k): + # Validate the input shapes. + assert_is_matrix(x) + assert_is_vector(indices) + assert_is_vector(bins) + assert_equal(indices.shape[0], x.shape[0] * top_k) + + if weights is not None: + assert_equal(weights.shape[0], x.shape[0] * top_k) + + num_experts = bins.shape[0] + out = torch.zeros((num_experts, expert_capacity, x.shape[1]), dtype=x.dtype, device=x.device) + + _binned_copy[(num_experts, expert_capacity)]( + x, + out, + num_experts, + expert_capacity, + indices, + weights, + bins, + NUM_COLUMNS=x.shape[1], + A_TO_B=True, + TOP_K=top_k, + SCALE=weights is not None, + ) + return out + + +def binned_scatter(x, indices, weights, bins, top_k): + # Validate the input shapes. + assert_is_tensor(x, 3) + assert_is_vector(indices) + assert_is_vector(bins) + assert_equal(bins.shape[0], x.shape[0]) + + if weights is not None: + assert_equal(indices.shape[0], weights.shape[0]) + + num_experts, expert_capacity, hidden_size = x.shape + tokens = indices.shape[0] // top_k + out = torch.zeros((tokens, top_k, hidden_size), dtype=x.dtype, device=x.device) + _binned_copy[(num_experts, expert_capacity)]( + out, + x, + num_experts, + expert_capacity, + indices, + weights, + bins, + NUM_COLUMNS=hidden_size, + A_TO_B=False, + TOP_K=top_k, + SCALE=weights is not None, + ) + + # Reduce along the top-k dimension, if needed. + return out.sum(dim=1) if top_k > 1 else out.view(tokens, hidden_size) + + +# a: (tokens, hidden_size), real. +# b: (num_experts, expert_capacity, num_columns), real. +# indices: (tokens * top_k), integer. +# weights: (tokens * top_k), real. +# bins: (num_experts), integer. +@triton.autotune( + configs=[ + triton.Config({'BLOCK_X': 64}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=2), + triton.Config({'BLOCK_X': 256}, num_warps=2), + triton.Config({'BLOCK_X': 128}, num_warps=4), + triton.Config({'BLOCK_X': 256}, num_warps=4), + ], + key=['NUM_COLUMNS'], +) +@triton.jit +def _binned_copy_wgrad( + x, + grad, + wgrad, + num_experts, + expert_capacity, + indices, + bins, + NUM_COLUMNS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_X: tl.constexpr, +): + # Load our indices into the output. + expert_idx = tl.program_id(0) + entry_idx = tl.program_id(1) + + # Calculate our offset into the output. + index_x = expert_idx * expert_capacity + entry_idx + + # Load the index bounds for our bin and calculate + # the number of tokens assigned to our expert. + start = 0 + if expert_idx > 0: + start = tl.load(bins + expert_idx - 1) + end = tl.load(bins + expert_idx) + num_tokens = end - start + + # Calculate our offset into the input. If we don't + # have an input exit early. + if entry_idx >= num_tokens: + return + index_out = tl.load(indices + start + entry_idx) + + # Offset the input and output pointers. + wgrad += index_out + grad += tl.multiple_of((index_out // TOP_K) * NUM_COLUMNS, NUM_COLUMNS) + x += tl.multiple_of(index_x * NUM_COLUMNS, NUM_COLUMNS) + offsets = tl.max_contiguous(tl.arange(0, BLOCK_X), BLOCK_X) + + acc = tl.zeros((BLOCK_X,), dtype=tl.float32) + iterations = tl.cdiv(NUM_COLUMNS, BLOCK_X) + for _ in range(iterations): + mask = offsets < NUM_COLUMNS + data = tl.load(x + offsets, mask=mask).to(tl.float32) + scale = tl.load(grad + offsets, mask=mask).to(tl.float32) + acc += data * scale + offsets += BLOCK_X + + # Reduce to get the final result and store. + out = tl.sum(acc).to(wgrad.dtype.element_ty) + tl.store(wgrad, out) + + +def binned_scatter_wgrad(x, grad, indices, bins, top_k): + # Validate the input shapes. + assert_is_tensor(x, 3) + assert_is_matrix(grad) + assert_is_vector(indices) + assert_is_vector(bins) + assert_equal(bins.shape[0], x.shape[0]) + + num_experts, expert_capacity, hidden_size = x.shape + tokens = indices.shape[0] // top_k + out = torch.zeros((tokens * top_k), dtype=x.dtype, device=x.device) + _binned_copy_wgrad[(num_experts, expert_capacity)]( + x, + grad, + out, + num_experts, + expert_capacity, + indices, + bins, + NUM_COLUMNS=hidden_size, + TOP_K=top_k, + ) + return out diff --git a/build/torch214-cxx11-cu132-x86_64-linux/benchmark_util.py b/build/torch214-cxx11-cu132-x86_64-linux/benchmark_util.py new file mode 100644 index 0000000000000000000000000000000000000000..02612d95e3ead1175a596e2878fa34b5bf85ad6f --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/benchmark_util.py @@ -0,0 +1,35 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import torch + + +def log_benchmark(name, arguments, time, std): + print('=' * 60) + print(f'{name} Benchmark') + print('Benchmark Parameters:') + for (key, value) in arguments.items(): + print(f'{key} = {value}') + print('Results:') + print('mean time = {:.3f}ms, std time = {:.3f}ms'.format(time, std)) + print('=' * 60) + + +def benchmark_function(fn, iterations=100, warmup=10): + # Warmup iterations. + for _ in range(warmup): + fn() + + times = [] + for i in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + start.record() + fn() + end.record() + + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + return np.mean(times), np.std(times) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/cpu_fused_moe.py b/build/torch214-cxx11-cu132-x86_64-linux/cpu_fused_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..7a29e81d6e9c57f64b0a78ac1c0828e45fd9d855 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/cpu_fused_moe.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# MegaBlocks CPU Fused MoE Implementation +# +# This is a pure Python/PyTorch implementation for CPU. +# For better performance, consider using the C++ kernel implementation. +# +import torch +import torch.nn.functional as F + + +def swigluoai_activation(gate: torch.Tensor, up: torch.Tensor, + alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor: + """ + SwigluOAI activation function used in GptOss models. + + Formula: + gate = clamp(gate, max=limit) + up = clamp(up, -limit, limit) + glu = gate * sigmoid(gate * alpha) + output = (up + 1) * glu + + Args: + gate: Gate tensor from gate projection + up: Up tensor from up projection + alpha: Scaling factor for sigmoid (default: 1.702) + limit: Clamp limit (default: 7.0) + + Returns: + Activated tensor + """ + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + glu = gate * torch.sigmoid(gate * alpha) + return (up + 1) * glu + + +def silu_and_mul_activation(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """ + SiLU (Swish) activation with element-wise multiplication. + + Formula: + output = silu(gate) * up + + Args: + gate: Gate tensor + up: Up tensor + + Returns: + Activated tensor + """ + return F.silu(gate) * up + + +def route_tokens_cpu( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor | None, + moe_top_k: int, + moe_num_experts: int, + moe_normalize_expert_weights: int | None = None, +) -> tuple: + """ + Route tokens to experts and compute expert weights and indices (CPU version). + + Args: + x: Input tensor [batch, seq, hidden] or [tokens, hidden] + router_weight: Router weight [num_experts, hidden] + router_bias: Router bias [num_experts] or None + moe_top_k: Number of experts per token + moe_num_experts: Total number of experts + moe_normalize_expert_weights: Normalization order or None + + Returns: + Tuple of (logits, expert_weights, expert_indices) + """ + x_flat = x.view(-1, x.shape[-1]) + logits = F.linear(x_flat, router_weight, router_bias) + + if moe_top_k == 1: + expert_weights, expert_indices = logits.max(dim=-1, keepdim=True) + else: + expert_weights, expert_indices = torch.topk(logits, moe_top_k, dim=-1) + + expert_weights = expert_weights.softmax(dim=-1) + + if moe_normalize_expert_weights is not None: + expert_weights = expert_weights / torch.norm( + expert_weights, + p=moe_normalize_expert_weights, + dim=-1, + keepdim=True, + ) + + return logits, expert_weights, expert_indices + + +def cpu_fused_moe( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + activation: str = "silu", + alpha: float = 1.702, + limit: float = 7.0, + is_interleaved: bool = True, +) -> torch.Tensor: + """ + CPU Fused MoE using PyTorch operations. + + This implementation processes all experts in parallel using batched operations + instead of sequential for loops, which is more efficient on CPU. + + Args: + hidden_states: [num_tokens, hidden_size] + w1: [num_experts, hidden_size, 2*inter_size] - gate_up_proj weights + w2: [num_experts, inter_size, hidden_size] - down_proj weights + topk_weights: [num_tokens, topk] - routing weights + topk_ids: [num_tokens, topk] - expert indices + w1_bias: [num_experts, 2*inter_size] or None + w2_bias: [num_experts, hidden_size] or None + activation: "silu" or "swigluoai" + alpha: swigluoai alpha parameter + limit: swigluoai limit parameter + is_interleaved: whether gate_up is interleaved [g0,u0,g1,u1,...] (True for GptOss) + + Returns: + output: [num_tokens, hidden_size] + """ + num_tokens, hidden_size = hidden_states.shape + num_experts = w1.shape[0] + inter_size = w2.shape[1] + topk = topk_weights.shape[1] + + # Initialize output + output = torch.zeros_like(hidden_states) + + # Build expert mask: which tokens go to which expert + # expert_mask[expert_id] contains indices of (token_idx, topk_pos) pairs + for expert_idx in range(num_experts): + # Find tokens assigned to this expert + # mask shape: [num_tokens, topk], True where topk_ids == expert_idx + mask = (topk_ids == expert_idx) + + if not mask.any(): + continue + + # Get token indices and topk positions + token_indices, topk_positions = torch.where(mask) + + if len(token_indices) == 0: + continue + + # Gather input tokens for this expert + # current_hidden: [num_selected_tokens, hidden_size] + current_hidden = hidden_states[token_indices] + + # Get weights for this expert + # w1[expert_idx]: [hidden_size, 2*inter_size] + # w2[expert_idx]: [inter_size, hidden_size] + expert_w1 = w1[expert_idx] # [hidden_size, 2*inter_size] + expert_w2 = w2[expert_idx] # [inter_size, hidden_size] + + # First projection: hidden @ w1 -> [num_selected, 2*inter_size] + gate_up = current_hidden @ expert_w1 + + # Add bias if present + if w1_bias is not None: + gate_up = gate_up + w1_bias[expert_idx] + + # Split gate and up projections + if is_interleaved: + # GptOss uses interleaved layout: [g0, u0, g1, u1, ...] + gate = gate_up[..., ::2] # [num_selected, inter_size] + up = gate_up[..., 1::2] # [num_selected, inter_size] + else: + # Standard layout: [gate_all, up_all] + gate = gate_up[..., :inter_size] + up = gate_up[..., inter_size:] + + # Apply activation + if activation == "swigluoai": + activated = swigluoai_activation(gate, up, alpha, limit) + else: # silu + activated = silu_and_mul_activation(gate, up) + + # Second projection: activated @ w2 -> [num_selected, hidden_size] + expert_out = activated @ expert_w2 + + # Add bias if present + if w2_bias is not None: + expert_out = expert_out + w2_bias[expert_idx] + + # Apply routing weights and accumulate + # weights shape: [num_selected] + weights = topk_weights[token_indices, topk_positions].unsqueeze(-1) + weighted_out = expert_out * weights + + # Accumulate to output + output.index_add_(0, token_indices, weighted_out.to(output.dtype)) + + return output + + +class MegaBlocksMoeMLP(torch.nn.Module): + """ + CPU MoE MLP module that can be used as a drop-in replacement for + the transformers GptOssMLP when using @use_kernel_forward_from_hub. + """ + can_torch_compile: bool = True + + def forward(self, x: torch.Tensor) -> tuple: + """ + Forward pass through the MoE layer. + + Args: + x: Input tensor of shape [batch_size, seq_len, hidden_size] or [tokens, hidden_size] + + Returns: + Tuple of (output, expert_weights) where: + - output: Tensor of same shape as input + - expert_weights: Expert weights for each token [tokens, top_k] + """ + # Get MoE parameters from the wrapped modules + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + moe_normalize_expert_weights = getattr( + self.experts, "normalize_expert_weights", None + ) + + # Detect activation type + if hasattr(self.experts, "alpha") and hasattr(self.experts, "limit"): + activation = "swigluoai" + alpha = self.experts.alpha + limit = self.experts.limit + else: + activation = getattr(self.experts, "activation", "silu") + alpha = 1.702 + limit = 7.0 + + # Get weight tensors + if hasattr(self.experts, "gate_up_proj"): + w1 = self.experts.gate_up_proj + is_interleaved = True # GptOss uses interleaved layout + elif hasattr(self.experts, "w1"): + w1 = self.experts.w1 + w3 = getattr(self.experts, "w3", None) + if w3 is not None: + w1 = torch.cat([w1, w3], dim=-1) + is_interleaved = False + else: + raise AttributeError("experts module must have 'gate_up_proj' or 'w1' attribute") + + if hasattr(self.experts, "down_proj"): + w2 = self.experts.down_proj + elif hasattr(self.experts, "w2"): + w2 = self.experts.w2 + else: + raise AttributeError("experts module must have 'down_proj' or 'w2' attribute") + + # Get optional bias tensors + w1_bias = getattr(self.experts, "gate_up_proj_bias", None) + w2_bias = getattr(self.experts, "down_proj_bias", None) + + # Store original shape + in_shape = x.size() + + # Route tokens to experts + logits, expert_weights, expert_indices = route_tokens_cpu( + x, + self.router.weight, + getattr(self.router, "bias", None), + moe_top_k, + moe_num_experts, + moe_normalize_expert_weights, + ) + + # Reshape input for fused MoE + x_flat = x.view(-1, x.shape[-1]) + + # Call CPU fused MoE + output = cpu_fused_moe( + hidden_states=x_flat, + w1=w1, + w2=w2, + topk_weights=expert_weights, + topk_ids=expert_indices, + w1_bias=w1_bias, + w2_bias=w2_bias, + activation=activation, + alpha=alpha, + limit=limit, + is_interleaved=is_interleaved, + ) + + # Restore original shape + output = output.view(in_shape) + + return output, expert_weights + + +# Export classes and functions +__all__ = [ + "MegaBlocksMoeMLP", + "cpu_fused_moe", + "route_tokens_cpu", + "swigluoai_activation", + "silu_and_mul_activation", +] diff --git a/build/torch214-cxx11-cu132-x86_64-linux/cpu_moe_cpp.py b/build/torch214-cxx11-cu132-x86_64-linux/cpu_moe_cpp.py new file mode 100644 index 0000000000000000000000000000000000000000..073ff66d24ce348fbb5ed19c9027fadd3f7a9c61 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/cpu_moe_cpp.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +# MegaBlocks C++ Optimized CPU MoE + +""" +C++ accelerated MoE with brgemm optimization for Intel AMX. +Direct replacement for cpu_fused_moe.MegaBlocksMoeMLP with better performance. +""" + +import torch +from typing import Optional +from .cpu_fused_moe import route_tokens_cpu +from ._ops import ops + + +def _to_local_tensor(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + """Convert DTensor to local torch.Tensor if needed for custom ops compatibility.""" + if tensor is None: + return None + # Check if it's a DTensor by looking for the to_local() method + if hasattr(tensor, "to_local"): + return tensor.to_local() + return tensor + + +def fused_moe_cpp( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + inplace: bool = False, + use_int8_w8a8: bool = False, + use_fp8_w8a16: bool = False, + use_mxfp4: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + block_size: Optional[list] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + w1_bias: Optional[torch.Tensor] = None, + w2_bias: Optional[torch.Tensor] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + is_vnni: bool = False, +) -> torch.Tensor: + """ + C++ Fused MoE with brgemm optimization (sglang compatible interface). + + Uses at::native::cpublas::brgemm for efficient batch GEMM on Intel CPUs. + Supports both silu_and_mul (standard SwiGLU) and swigluoai (GptOss) activations. + + Args: + hidden_states: Input tensor [M, K] + w1: Gate and up projections [E, 2N, K] + w2: Down projection [E, K, N] + topk_weights: Expert weights [M, topk] + topk_ids: Expert indices [M, topk] + inplace: Whether to use hidden_states as output + use_int8_w8a8: Use int8 quantization + use_fp8_w8a16: Use fp8 quantization + use_mxfp4: Use mxfp4 quantization + w1_scale, w2_scale: Quantization scales + block_size: Block size for fp8 + a1_scale, a2_scale: Activation scales + w1_bias, w2_bias: Optional biases + alpha: swigluoai alpha parameter (set to enable swiglu) + limit: swigluoai limit parameter (set to enable swiglu) + is_vnni: Whether w1/w2 are already in VNNI packed format + """ + # MXFP4/FP8 kernels only support bf16, convert if needed + orig_dtype = hidden_states.dtype + need_convert = ((use_mxfp4 or use_fp8_w8a16) and orig_dtype != torch.bfloat16) or orig_dtype == torch.float32 + if need_convert: + hidden_states = hidden_states.to(torch.bfloat16) + + # bias must match hidden_states dtype + if w1_bias is not None: + w1_bias = w1_bias.to(hidden_states.dtype) + if w2_bias is not None: + w2_bias = w2_bias.to(hidden_states.dtype) + + # Convert DTensor to local tensor for custom ops compatibility (TP mode) + hidden_states = _to_local_tensor(hidden_states) + w1 = _to_local_tensor(w1) + w2 = _to_local_tensor(w2) + topk_weights = _to_local_tensor(topk_weights) + topk_ids = _to_local_tensor(topk_ids) + w1_scale = _to_local_tensor(w1_scale) + w2_scale = _to_local_tensor(w2_scale) + a1_scale = _to_local_tensor(a1_scale) + a2_scale = _to_local_tensor(a2_scale) + w1_bias = _to_local_tensor(w1_bias) + w2_bias = _to_local_tensor(w2_bias) + + output = ops.fused_experts( + hidden_states, w1, w2, topk_weights, topk_ids, + inplace, use_int8_w8a8, use_fp8_w8a16, use_mxfp4, + w1_scale, w2_scale, block_size, a1_scale, a2_scale, + w1_bias, w2_bias, alpha, limit, is_vnni + ) + + # Convert back to original dtype if needed + if need_convert: + output = output.to(orig_dtype) + return output + + +class CPUMegaBlocksMoeMLP(torch.nn.Module): + """ + C++ optimized MoE MLP using brgemm. + Drop-in replacement for cpu_fused_moe.MegaBlocksMoeMLP with better performance. + + Usage in transformers: + # Will be used via @use_kernel_forward_from_hub decorator + """ + can_torch_compile: bool = True + + def forward(self, x: torch.Tensor) -> tuple: + """ + Forward pass through the MoE layer using C++ kernel. + + Args: + x: Input tensor [batch_size, seq_len, hidden_size] + + Returns: + Tuple of (output, expert_weights) + """ + # Optimization for GPT-OSS model + if getattr(self, "use_mxfp4", None) is None: + self.use_mxfp4 = False + + w1_scale = None + w2_scale = None + + if ( + not getattr(self, "packed_scales", False) + and hasattr(self.experts, "gate_up_proj") + and getattr(self.experts, "gate_up_proj_precision_config", None) is not None + ): + # convert scales + data_1 = ops.convert_scale_packed(self.experts.gate_up_proj_precision_config.weight_scale.data.transpose(-1, -2).contiguous()) + data_2 = ops.convert_scale_packed(self.experts.down_proj_precision_config.weight_scale.data.transpose(-1, -2).contiguous()) + self.experts.gate_up_proj_precision_config.weight_scale.storage.data = data_1 + self.experts.down_proj_precision_config.weight_scale.storage.data = data_2 + self.packed_scales = True + self.use_mxfp4 = True + + if not getattr(self, "packed_weight", False) and hasattr( + self.experts, "gate_up_proj" + ): + # convert weights + data_1 = self.experts.gate_up_proj.data.transpose(-1, -2).contiguous() + data_2 = self.experts.down_proj.data.transpose(-1, -2).contiguous() + if self.use_mxfp4: + self.experts.gate_up_proj.storage.data = ops.convert_weight_packed(data_1) + self.experts.down_proj.storage.data = ops.convert_weight_packed(data_2) + else: + # convert_weight_packed only supports bfloat16, float16, int8, fp8_e4m3 or uint8(mxfp4 or int4). + data_1 = data_1.to(torch.bfloat16) if data_1.dtype == torch.float32 else data_1 + data_2 = data_2.to(torch.bfloat16) if data_2.dtype == torch.float32 else data_2 + self.experts.gate_up_proj.data = ops.convert_weight_packed(data_1) + self.experts.down_proj.data = ops.convert_weight_packed(data_2) + + # C++ kernel does not support float32. + dtype = torch.bfloat16 if x.dtype == torch.float32 else x.dtype + if getattr(self.experts, "gate_up_proj_bias", None) is not None: + self.experts.gate_up_proj_bias.data = self.experts.gate_up_proj_bias.data.to(dtype) + if getattr(self.experts, "down_proj_bias", None) is not None: + self.experts.down_proj_bias.data = self.experts.down_proj_bias.data.to(dtype) + + self.packed_weight = True + + # Get MoE parameters + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None) + + # Detect activation type + if hasattr(self.experts, "alpha") and hasattr(self.experts, "limit"): + activation = "swigluoai" + alpha = self.experts.alpha + limit = self.experts.limit + else: + activation = getattr(self.experts, "activation", "silu") + alpha = 1.702 + limit = 7.0 + + # Get weight tensors + if hasattr(self.experts, "gate_up_proj"): + w1 = self.experts.gate_up_proj + elif hasattr(self.experts, "w1"): + w1 = self.experts.w1 + w3 = getattr(self.experts, "w3", None) + if w3 is not None: + w1 = torch.cat([w1, w3], dim=-1) + else: + raise AttributeError("experts module must have 'gate_up_proj' or 'w1' attribute") + + if hasattr(self.experts, "down_proj"): + w2 = self.experts.down_proj + elif hasattr(self.experts, "w2"): + w2 = self.experts.w2 + else: + raise AttributeError("experts module must have 'down_proj' or 'w2' attribute") + + # Get optional bias tensors + w1_bias = getattr(self.experts, "gate_up_proj_bias", None) + w2_bias = getattr(self.experts, "down_proj_bias", None) + w1_bias = w1_bias if w1_bias is None else w1_bias.data + w2_bias = w2_bias if w2_bias is None else w2_bias.data + + if self.use_mxfp4: + w1_scale = self.experts.gate_up_proj_precision_config.weight_scale.data + w2_scale = self.experts.down_proj_precision_config.weight_scale.data + + # Store original shape + in_shape = x.size() + + # Route tokens to experts (Python implementation is fast enough) + logits, expert_weights, expert_indices = route_tokens_cpu( + x, + self.router.weight, + getattr(self.router, "bias", None), + moe_top_k, + moe_num_experts, + moe_normalize_expert_weights, + ) + + # Flatten input + x_flat = x.view(-1, x.shape[-1]) + + # Determine alpha/limit for swiglu activation + use_alpha = alpha if activation == "swigluoai" else None + use_limit = limit if activation == "swigluoai" else None + + # Call C++ optimized kernel + output = fused_moe_cpp( + hidden_states=x_flat, + w1=w1.data, + w2=w2.data, + topk_weights=expert_weights, + topk_ids=expert_indices.to(torch.int32), + inplace=False, + use_int8_w8a8=False, + use_fp8_w8a16=False, + use_mxfp4=self.use_mxfp4, + w1_scale=w1_scale, + w2_scale=w2_scale, + block_size=None, + a1_scale=None, + a2_scale=None, + w1_bias=w1_bias, + w2_bias=w2_bias, + alpha=use_alpha, + limit=use_limit, + is_vnni=getattr(self, "packed_weight", False), + ) + + # Restore original shape + output = output.view(in_shape) + + return output, expert_weights + + +__all__ = ["fused_moe_cpp", "MegaBlocksMoeMLP"] diff --git a/build/torch214-cxx11-cu132-x86_64-linux/grouped_gemm/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/grouped_gemm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b91c8308f0c24f4c4171b6e4f15b6f76dabf295a --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/grouped_gemm/__init__.py @@ -0,0 +1,2 @@ +from . import ops +from . import backend diff --git a/build/torch214-cxx11-cu132-x86_64-linux/grouped_gemm/backend.py b/build/torch214-cxx11-cu132-x86_64-linux/grouped_gemm/backend.py new file mode 100644 index 0000000000000000000000000000000000000000..c7ef28ced79c830dae934177f059c1f4ddc24aad --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/grouped_gemm/backend.py @@ -0,0 +1,33 @@ +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# On ROCm there is no CUTLASS grouped GEMM; dispatch to the vendored AITER +# Triton kernels instead. On CUDA we use the compiled CUTLASS `gmm` op. +_IS_ROCM = torch.version.hip is not None + +if _IS_ROCM: + from .._grouped_gemm_triton import adapter as backend +else: + # We import the backend operations from the megablocks package as + # grouped_gemm is vendored in megablocks in this repository. + from .._ops import ops as backend # type: ignore + +def _allocate_output(a, b, batch_sizes, trans_a, trans_b): + assert not (trans_a and trans_b) + assert batch_sizes.ndim == 1, "Expected 1d tensor for batch_sizes" + assert a.ndim == 2, "Expected 2d tensor for 'a'" + assert b.ndim == (2 if trans_a else 3) + + shape = ( + (batch_sizes.shape[0], a.shape[1], b.shape[1]) + if trans_a else + (a.shape[0], (b.shape[1] if trans_b else b.shape[2])) + ) + return torch.empty(*shape, device=a.device, dtype=a.dtype) + +def gmm(a, b, batch_sizes, trans_a=False, trans_b=False, c=None): + if c is None: + c = _allocate_output(a, b, batch_sizes, trans_a, trans_b) + backend.gmm(a, b, c, batch_sizes, trans_a, trans_b) + return c diff --git a/build/torch214-cxx11-cu132-x86_64-linux/grouped_gemm/ops.py b/build/torch214-cxx11-cu132-x86_64-linux/grouped_gemm/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..4b30dd14e23837ea3b12334f4e31337ed9ad2b69 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/grouped_gemm/ops.py @@ -0,0 +1,33 @@ +from . import backend +import torch + + +class GroupedGemm(torch.autograd.Function): + + @staticmethod + def forward(ctx, a, b, batch_sizes, trans_b): + ctx.save_for_backward(a, b, batch_sizes) + ctx.trans_b = trans_b + return backend.gmm(a, b, batch_sizes, trans_a=False, trans_b=trans_b) + + @staticmethod + def backward(ctx, grad): + grad = grad.contiguous() + a, b, batch_sizes = ctx.saved_tensors + trans_b = ctx.trans_b + + agrad = None + if ctx.needs_input_grad[0]: + agrad = backend.gmm( + grad, b, batch_sizes, trans_a=False, trans_b=not trans_b) + + bgrad = None + if ctx.needs_input_grad[1]: + lhs, rhs = (grad, a) if trans_b else (a, grad) + bgrad = backend.gmm( + lhs, rhs, batch_sizes, trans_a=True, trans_b=False) + return agrad, bgrad, None, None + + +def gmm(a, b, batch_sizes, trans_b=False): + return GroupedGemm.apply(a, b, batch_sizes, trans_b) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/grouped_gemm_util.py b/build/torch214-cxx11-cu132-x86_64-linux/grouped_gemm_util.py new file mode 100644 index 0000000000000000000000000000000000000000..1d6d49fc46b0a57ad46e4179df3cc1ac2a24f7ae --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/grouped_gemm_util.py @@ -0,0 +1,31 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +import warnings + +_grouped_gemm_is_available: bool = False +try: + # import grouped_gemm + pass + _grouped_gemm_is_available = True +except ImportError as error: + warnings.warn('Grouped GEMM not available.') + + +def grouped_gemm_is_available(): + return _grouped_gemm_is_available + + +def assert_grouped_gemm_is_available(): + msg = ( + 'Grouped GEMM not available. Please run ' + '`pip install git+https://github.com/tgale96/grouped_gemm@main`.', + ) + assert _grouped_gemm_is_available, msg + + +# backend = grouped_gemm.backend if grouped_gemm_is_available() else None +# ops = grouped_gemm.ops if grouped_gemm_is_available() else None + + +from .grouped_gemm import backend as ops +from .grouped_gemm import ops as backend diff --git a/build/torch214-cxx11-cu132-x86_64-linux/layers.py b/build/torch214-cxx11-cu132-x86_64-linux/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..7c91edcd9e2d1a4ef9eac90217ff481f08ab1886 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/layers.py @@ -0,0 +1,1232 @@ +import torch +import torch.distributed as dist + +from typing import Optional, Any, TYPE_CHECKING + +from . import _layers +from . import ops + +# Conditional import for meta kernel registration +if TYPE_CHECKING: + + def register_fake(fn): + return lambda name: fn + +else: + try: + from torch.library import register_fake + except ImportError: + try: + from torch.library import impl_abstract as register_fake + except ImportError: + # Fallback for older PyTorch versions + def register_fake(op_name): + def decorator(fn): + return fn + + return decorator + + +# Meta kernel implementations for torch.compile compatibility +def _install_meta_kernels(): + """Install meta kernels for existing MegaBlocks operations""" + + # Create wrapper functions that check for compilation and return meta tensors + + # Patch ops.sort + if hasattr(ops, "sort"): + original_sort = ops.sort + + def sort_with_meta(x, end_bit=None): + if torch.compiler.is_compiling(): + # print("Using meta kernel for sort") + # Meta implementation - return tensors with correct shape/dtype/device + return torch.empty_like(x), torch.empty_like(x) + # print("Using original sort kernel") + return original_sort(x, end_bit) + + ops.sort = sort_with_meta + + # Patch ops.histogram + if hasattr(ops, "histogram"): + original_histogram = ops.histogram + + def histogram_with_meta(x, max_val): + if torch.compiler.is_compiling(): + # Meta implementation + return torch.empty((max_val,), dtype=torch.int32, device=x.device) + return original_histogram(x, max_val) + + ops.histogram = histogram_with_meta + + # Patch ops.inclusive_cumsum + if hasattr(ops, "inclusive_cumsum"): + original_inclusive_cumsum = ops.inclusive_cumsum + + def inclusive_cumsum_with_meta(x, dim): + if torch.compiler.is_compiling(): + # Meta implementation + return torch.empty_like(x) + return original_inclusive_cumsum(x, dim) + + ops.inclusive_cumsum = inclusive_cumsum_with_meta + + # Patch ops.binned_gather + if hasattr(ops, "binned_gather"): + original_binned_gather = ops.binned_gather + + def binned_gather_with_meta(x, indices, bins, bin_size, top_k): + if torch.compiler.is_compiling(): + # Meta implementation - output shape based on bin_size + if x.dim() >= 2: + hidden_size = x.size(-1) + return torch.empty( + (bin_size, x.size(1), hidden_size), + dtype=x.dtype, + device=x.device, + ) + else: + return torch.empty((bin_size,), dtype=x.dtype, device=x.device) + return original_binned_gather(x, indices, bins, bin_size, top_k) + + ops.binned_gather = binned_gather_with_meta + + # Patch ops.binned_scatter + if hasattr(ops, "binned_scatter"): + original_binned_scatter = ops.binned_scatter + + def binned_scatter_with_meta(x, indices, weights, bins, top_k): + if torch.compiler.is_compiling(): + # Meta implementation - typically reduces to 2D + if x.dim() >= 3: + return torch.empty( + (x.size(1), x.size(2)), dtype=x.dtype, device=x.device + ) + else: + return torch.empty_like(x) + return original_binned_scatter(x, indices, weights, bins, top_k) + + ops.binned_scatter = binned_scatter_with_meta + + # Patch ops.gather + if hasattr(ops, "gather"): + original_gather = ops.gather + + def gather_with_meta(x, indices, bin_ids, bins, top_k): + if torch.compiler.is_compiling(): + # Meta implementation + if x.dim() >= 2: + hidden_size = x.size(-1) + return torch.empty( + (indices.numel(), hidden_size), dtype=x.dtype, device=x.device + ) + else: + return torch.empty(indices.shape, dtype=x.dtype, device=x.device) + return original_gather(x, indices, bin_ids, bins, top_k) + + ops.gather = gather_with_meta + + # Patch ops.scatter + if hasattr(ops, "scatter"): + original_scatter = ops.scatter + + def scatter_with_meta(x, indices, bin_ids, weights, bins, top_k): + if torch.compiler.is_compiling(): + # Meta implementation - restore sequence shape + seq_len = ( + indices.size(0) // top_k + if indices.numel() > 0 and top_k > 0 + else x.size(0) + ) + if x.dim() >= 2: + return torch.empty( + (seq_len, x.size(-1)), dtype=x.dtype, device=x.device + ) + else: + return torch.empty((seq_len,), dtype=x.dtype, device=x.device) + return original_scatter(x, indices, bin_ids, weights, bins, top_k) + + ops.scatter = scatter_with_meta + + # Patch ops.replicate + if hasattr(ops, "replicate"): + original_replicate = ops.replicate + + def replicate_with_meta(x, bins, num_outputs): + if torch.compiler.is_compiling(): + # Meta implementation + return torch.empty( + (x.shape[0], num_outputs), dtype=x.dtype, device=x.device + ) + return original_replicate(x, bins, num_outputs) + + ops.replicate = replicate_with_meta + + # Patch ops.repeat (if it's a regular function) + if hasattr(ops, "repeat"): + original_repeat = ops.repeat + + def repeat_with_meta(x, repeats): + if torch.compiler.is_compiling(): + # Meta implementation + if isinstance(repeats, (tuple, list)): + new_shape = list(x.shape) + for i, rep in enumerate(repeats): + if i < len(new_shape): + new_shape[i] *= rep + return torch.empty(new_shape, dtype=x.dtype, device=x.device) + else: + new_shape = [x.size(0) * repeats] + list(x.shape[1:]) + return torch.empty(new_shape, dtype=x.dtype, device=x.device) + return original_repeat(x, repeats) + + ops.repeat = repeat_with_meta + + +# Install meta kernels on import +try: + _install_meta_kernels() +except Exception as e: + # If meta kernel installation fails, continue without them + # torch.compile may not work but the library will still function + import warnings + + warnings.warn( + f"Failed to install meta kernels for torch.compile support: {e}", UserWarning + ) + + +# Set the expert model parallel attributes on a tensor +def set_expert_model_parallel_attributes( + tensor: torch.Tensor, + is_parallel: bool, +): + assert not hasattr(tensor, "expert_model_parallel") + setattr(tensor, "expert_model_parallel", is_parallel) + + +# Get the expert model parallel attributes from a tensor +def expert_sharding_degree( + world_size: int, + moe_num_experts: int, +) -> int: + esd = min(world_size, moe_num_experts) + if (moe_num_experts % esd) != 0: + raise ValueError(f"Cannot shard {moe_num_experts} experts {esd} ways.") + return esd + + +# Calculate the hidden sharding degree based on world size and expert sharding degree +def hidden_sharding_degree( + world_size: int, + moe_num_experts: int, + ffn_hidden_size: int, +) -> int: + esd = expert_sharding_degree(world_size, moe_num_experts) + hsd = world_size // esd + if (ffn_hidden_size % hsd) != 0: + raise ValueError(f"Cannot shard {ffn_hidden_size} features {hsd} ways.") + if (esd * hsd) != world_size: + raise ValueError( + f"Invalid sharding. expert_sharding_degree ({esd}) * hidden_sharding_degree ({hsd}) != world_size ({world_size})." + ) + return hsd + + +# Calculate the number of experts per rank based on world size and expert sharding degree +def experts_per_rank( + moe_num_experts: int, + world_size: int, +) -> int: + return moe_num_experts // expert_sharding_degree(world_size, moe_num_experts) + + +# Calculate the number of features per rank based on ffn hidden size and hidden sharding degree +def features_per_rank( + ffn_hidden_size: int, world_size: int, moe_num_experts: int +) -> int: + return ffn_hidden_size // hidden_sharding_degree( + world_size, moe_num_experts, ffn_hidden_size + ) + + +# Apply jitter to the input tensor +def apply_jitter(x: torch.Tensor, moe_jitter_eps: float) -> torch.Tensor: + low = 1.0 - moe_jitter_eps + high = 1.0 + moe_jitter_eps + noise = torch.rand(x.size(), dtype=x.dtype, device=x.device) + return x * (low + noise * (high - low)) + + +# Compute the top-k scores from the logits +def compute_top_k(scores: torch.Tensor, moe_top_k: int): + if moe_top_k == 1: + return scores.max(dim=-1, keepdim=True) + return torch.topk(scores, moe_top_k, dim=-1) + + +# Route tokens to experts and compute expert weights and indices +def route_tokens( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + moe_top_k: int, + moe_num_experts: int, + moe_jitter_eps: float = None, + moe_normalize_expert_weights: int = None, + uniform_expert_assignment: bool = False, + training: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if training and moe_jitter_eps is not None: + x = apply_jitter(x, moe_jitter_eps) + + x_flat = x.view(-1, x.shape[-1]) + logits = torch.nn.functional.linear(x_flat, router_weight, router_bias) + expert_weights, expert_indices = compute_top_k(logits, moe_top_k) + expert_weights = expert_weights.softmax(dim=-1) + if moe_normalize_expert_weights is not None: + expert_weights = expert_weights / torch.norm( + expert_weights, + p=moe_normalize_expert_weights, + dim=-1, + keepdim=True, + ) + if uniform_expert_assignment: + expert_indices = _layers.router._uniform_expert_assignment( + expert_indices, + moe_num_experts, + ) + + return logits, expert_weights, expert_indices + + +# Scale the gradient of the weights +def scale_grad( + w: torch.Tensor, + gradient_scale: Optional[float] = None, +) -> torch.Tensor: + if gradient_scale is None: + return w + return _layers.mlp.scale_gradient(w, gradient_scale) + + +# Forward pass for the MLP layer +def mlp_forward( + x: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w1_bias: torch.Tensor, + w2_bias: torch.Tensor, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + limit: float = 7.0, +): + # Scale weights + w1 = scale_grad(w1, gradient_scale) + w2 = scale_grad(w2, gradient_scale) + w1_bias = scale_grad(w1_bias, gradient_scale) + w2_bias = scale_grad(w2_bias, gradient_scale) + + # Resolve dtensors + w1 = _layers.mlp.resolve_dtensor(w1) + w2 = _layers.mlp.resolve_dtensor(w2) + w1_bias = _layers.mlp.resolve_dtensor(w1_bias) + w2_bias = _layers.mlp.resolve_dtensor(w2_bias) + + # Forward pass + gate_up = torch.bmm(x, w1) + w1_bias[..., None, :] + gate, up = gate_up[..., ::2], gate_up[..., 1::2] + gate = gate.clamp(min=None, max=limit) + up = up.clamp(min=-limit, max=limit) + glu = gate * torch.sigmoid(gate * alpha) + next_states = torch.bmm(((up + 1) * glu), w2) + next_states += w2_bias[..., None, :] + return next_states + +# Shared expert MLP forward pass +def shared_mlp_forward( + x: torch.Tensor, + up_proj_weight: torch.Tensor, + down_proj_weight: torch.Tensor, + up_proj_bias: Optional[torch.Tensor] = None, + down_proj_bias: Optional[torch.Tensor] = None, + activation_fn: Optional[Any] = None, + gradient_scale: Optional[float] = None, +) -> torch.Tensor: + # Default activation function + if activation_fn is None: + activation_fn = torch.nn.functional.gelu + + # Scale weights + up_proj_weight = scale_grad(up_proj_weight, gradient_scale) + down_proj_weight = scale_grad(down_proj_weight, gradient_scale) + if up_proj_bias is not None: + up_proj_bias = scale_grad(up_proj_bias, gradient_scale) + if down_proj_bias is not None: + down_proj_bias = scale_grad(down_proj_bias, gradient_scale) + + # Resolve dtensors + up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight) + down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight) + if up_proj_bias is not None: + up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias) + if down_proj_bias is not None: + down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias) + + # Up projection + x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias) + + # Activation + x = activation_fn(x) + + # Down projection + x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias) + + return x + + +# Combine outputs from shared expert and regular experts +def combine_expert_shared_outputs( + shared_expert_out: torch.Tensor, + expert_out: torch.Tensor, + shared_expert_weighted_sum: bool = False, + moe_top_k: int = 1, +) -> torch.Tensor: + if shared_expert_weighted_sum: + # Weighted sum based on number of experts used + total_experts = moe_top_k + 1 + shared_weight = 1.0 / total_experts + expert_weight = moe_top_k / total_experts + return shared_expert_out * shared_weight + expert_out * expert_weight + else: + # Simple addition + return shared_expert_out + expert_out + + +# Global variable to store load balancing loss +_LOAD_BALANCING_LOSS = [] + + +def save_load_balancing_loss(loss): + global _LOAD_BALANCING_LOSS + _LOAD_BALANCING_LOSS.append(loss) + + +def get_load_balancing_loss(): + global _LOAD_BALANCING_LOSS + return _LOAD_BALANCING_LOSS + + +def clear_load_balancing_loss(): + global _LOAD_BALANCING_LOSS + _LOAD_BALANCING_LOSS.clear() + + +def batched_load_balancing_loss(args): + if args.moe_loss_weight == 0: + return 0.0 + + tokens_per_expert, expert_scores = zip(*get_load_balancing_loss()) + num_layers_per_pipeline_stage = args.num_layers // args.pipeline_model_parallel_size + if args.num_layers_per_virtual_pipeline_stage is not None: + num_layers_per_pipeline_stage = args.num_layers_per_virtual_pipeline_stage + + if len(tokens_per_expert) != num_layers_per_pipeline_stage: + raise ValueError( + f"Expected {num_layers_per_pipeline_stage} token_per_experts " + f"but found {len(tokens_per_expert)}.\nnum_layers = " + f"{args.num_layers}\npipeline_model_parallel_size = " + f"{args.pipeline_model_parallel_size}\n" + "num_layers_per_virtual_pipeline_stage" + f" = {args.num_layers_per_virtual_pipeline_stage}", + ) + if len(expert_scores) != num_layers_per_pipeline_stage: + raise ValueError( + f"Expected {num_layers_per_pipeline_stage} expert_scores " + f"but found {len(tokens_per_expert)}.\nnum_layers = " + f"{args.num_layers}\npipeline_model_parallel_size = " + f"{args.pipeline_model_parallel_size}\n" + "num_layers_per_virtual_pipeline_stage" + f" = {args.num_layers_per_virtual_pipeline_stage}", + ) + + # Verify the shape of the tokens_per_expert and expert_scores tensors. + assert all( + (x.ndim == 1 and x.numel() == args.moe_num_experts for x in tokens_per_expert) + ) + + tokens = expert_scores[0].shape[0] + assert all( + ( + ( + x.ndim == 2 + and x.shape[1] == args.moe_num_experts + and x.shape[0] == tokens + ) + for x in expert_scores + ) + ) + + # Concatenate the contributions of each layer and convert to + # the correct types and formats for the dot product. + expert_scores = torch.cat(expert_scores, dim=1) + if args.moe_lbl_in_fp32: + expert_scores = expert_scores.float() + if tokens != 0: + expert_scores = expert_scores.mean(dim=0) + else: + expert_scores = expert_scores.sum(dim=0) + tokens_per_expert = torch.cat(tokens_per_expert).to(expert_scores.dtype) + + expected_values = num_layers_per_pipeline_stage * args.moe_num_experts + assert tokens_per_expert.numel() == expected_values + assert expert_scores.numel() == expected_values + + # Calculate the total scale across all factors. + # + # loss_weight * num_experts / (num_layers * tokens * top_k) + scale_numerator = args.moe_num_experts * args.moe_loss_weight + scale_denominator = args.num_layers * tokens * args.moe_top_k + scale = scale_numerator / scale_denominator + return scale * torch.dot(tokens_per_expert, expert_scores) + + +# Calculate the expert capacity based on tokens, top_k, number of experts, +# expert parallel group, capacity factor, and whether expert model parallelism is used. +def expert_capacity( + tokens: int, + top_k: int, + num_experts: int, + expert_parallel_group: int, + moe_capacity_factor: float, + moe_expert_model_parallelism: bool, +) -> int: + world_size = ( + dist.get_world_size(expert_parallel_group) + if moe_expert_model_parallelism + else 1 + ) + + tokens_per_expert = top_k * tokens * world_size / num_experts + return int(moe_capacity_factor * tokens_per_expert) + + +def load_balancing_loss( + tokens_per_expert: torch.Tensor, + expert_scores: torch.Tensor, + top_k: int, + num_experts: int, +): + assert len(expert_scores.size()) == 2 + tokens, num_experts = expert_scores.size() + assert num_experts == num_experts + assert len(tokens_per_expert.size()) == 1 + (num_experts,) = tokens_per_expert.size() + assert num_experts == num_experts + scale = num_experts / (tokens * top_k) + return scale * torch.dot( + tokens_per_expert.to(expert_scores.dtype), + expert_scores.mean(dim=0), + ) + + +def indices_and_bins( + top_expert: torch.Tensor, + sort_end_bit: int, + num_experts: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + top_expert = top_expert.int() + + # Ensure contiguous memory layout + top_expert = top_expert.contiguous() + + # Ensure CUB knows which device to use + with torch.cuda.device(top_expert.device): + output = ops.sort(top_expert, sort_end_bit) + bin_ids, indices = output + tokens_per_expert = ops.histogram(top_expert, num_experts) + bins = ops.inclusive_cumsum(tokens_per_expert, 0) + + bins = bins.view(1) if not len(bins.size()) else bins + return indices, bin_ids, bins, tokens_per_expert + + +def expert_capacity_fn( + tokens: int, + top_k: int, + num_experts: int, + expert_parallel_group: torch.distributed.ProcessGroup, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = False, +) -> int: + world_size = ( + dist.get_world_size(expert_parallel_group) + if moe_expert_model_parallelism + else 1 + ) + tokens_per_expert = top_k * tokens * world_size / num_experts + return int(moe_capacity_factor * tokens_per_expert) + + +def permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capacity, + top_k, + w1, + w2, + w1_bias, + w2_bias, + gradient_scale, + alpha, +): + # Route tokens to experts + x = x.view(-1, x.shape[-1]) + + # Ensure CUB knows which device to use + with torch.cuda.device(x.device): + x = ops.binned_gather(x, indices, bins, expert_capacity, top_k) + + # Expert computation + x = mlp_forward(x, w1, w2, w1_bias, w2_bias, gradient_scale, alpha) + + # Ensure CUB knows which device to use + with torch.cuda.device(x.device): + # Route tokens back + out = ops.binned_scatter(x, indices, expert_weights, bins, top_k) + return out + + +def forward_once( + x: torch.Tensor, + expert_weights: torch.Tensor, + top_experts: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w1_bias: torch.Tensor, + w2_bias: torch.Tensor, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + sort_end_bit: int = 0, + top_k: int = 4, + num_experts: int = 128, + expert_parallel_group: int = None, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = False, + mlp_impl: Optional[str] = None, +): + # x: [sl, bs, hs] + # expert_weights: [sl * bs, top-k] + # top_experts: [sl * bs, top-k] + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + + with torch.no_grad(): + indices, bin_ids, bins, tokens_per_expert = indices_and_bins( + top_experts, sort_end_bit, num_experts + ) + + # Calculate expert capacity + sl, bs, _ = x.size() + + expert_capacity = expert_capacity_fn( + sl * bs, + top_k, + num_experts, + expert_parallel_group, + moe_capacity_factor, + moe_expert_model_parallelism, + ) + + if expert_capacity == 0: + expert_capacity = torch.max(tokens_per_expert).item() + + x = permute_and_compute( + x, + tokens_per_expert, + indices, + bin_ids, + expert_weights, + bins, + expert_capacity, + top_k, + w1, + w2, + w1_bias, + w2_bias, + gradient_scale, + alpha, + ) + return x, tokens_per_expert + + +def parallel_forward_once( + x: torch.Tensor, + expert_weights: torch.Tensor, + top_experts: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w1_bias: torch.Tensor, + w2_bias: torch.Tensor, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + sort_end_bit: int = 0, + top_k: int = 4, + num_experts: int = 128, + expert_parallel_group: torch.distributed.ProcessGroup = None, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = True, + hidden_size: int = 1152, + mlp_impl: Optional[str] = "grouped", +): + # Flatten inputs + expert_weights = expert_weights.flatten() + top_experts = top_experts.flatten() + + # TODO: remove debugging var + # my_rank = dist.get_rank(expert_parallel_group) if expert_parallel_group else 0 + + with torch.no_grad(): + # Step 1: Local permutation setup + indices, bin_ids, bins, tokens_per_expert = indices_and_bins( + top_experts, sort_end_bit, num_experts + ) + + # Calculate sharding parameters + world_size = dist.get_world_size(expert_parallel_group) + hidden_sharding_deg = hidden_sharding_degree( + world_size, num_experts, hidden_size + ) + experts_per_rank_val = experts_per_rank(num_experts, world_size) + + # Replicate token counts for hidden sharding + repeated_tokens_per_expert = ops.repeat( + tokens_per_expert, (hidden_sharding_deg,) + ) + + # Exchange token counts across devices + parallel_tokens_per_expert = torch.empty_like(repeated_tokens_per_expert) + + # Ensure CUB knows which device to use + tpe_handle = dist.all_to_all_single( + parallel_tokens_per_expert, + repeated_tokens_per_expert, + group=expert_parallel_group, + async_op=True, + ) + + # Step 2: Local permutation - group tokens by target device + x = x.view(-1, x.shape[-1]) # [sl * bs, hs] + x = ops.gather(x, indices, bin_ids, bins, top_k) + + # Step 3: Compute communication counts and exchange tokens + with torch.no_grad(): + tpe_handle.wait() + + # Reshape for per-device calculations + repeated_tokens_per_expert = repeated_tokens_per_expert.view( + world_size, experts_per_rank_val + ) + parallel_tokens_per_expert = parallel_tokens_per_expert.view( + world_size, experts_per_rank_val + ) + + # Calculate send/recv counts + send_counts = repeated_tokens_per_expert.cpu().sum(dim=-1).tolist() + # recv_counts = parallel_tokens_per_expert.cpu().sum(dim=-1).tolist() + parallel_tokens_per_expert_cpu = parallel_tokens_per_expert.cpu() + recv_counts = parallel_tokens_per_expert_cpu.sum(dim=-1).tolist() + tokens_received = sum(recv_counts) + + # Replicate for hidden sharding + x = ops.repeat(x, (hidden_sharding_deg, 1)) + + # Cross-device token exchange + parallel_x, parallel_x_handle = _layers.all_to_all.all_to_all( + x, recv_counts, send_counts, expert_parallel_group, async_op=True + ) + + with torch.no_grad(): + # Step 4: Setup for local expert computation + replicate_bins = ops.inclusive_cumsum(parallel_tokens_per_expert.flatten(), 0) + replicate_bins = ( + replicate_bins.view(1) if not len(replicate_bins.size()) else replicate_bins + ) + + # Create expert indices for received tokens + parallel_top_expert = torch.remainder( + torch.arange( + num_experts * hidden_sharding_deg, + dtype=torch.int32, + device=indices.device, + ), + experts_per_rank_val, + ) + parallel_top_expert = ops.replicate( + parallel_top_expert.unsqueeze(dim=0), + replicate_bins, + tokens_received, + ).flatten() + + # Sort tokens by expert assignment + parallel_bin_ids, parallel_indices = ops.sort( + parallel_top_expert, + sort_end_bit, + ) + + # Calculate bins for local experts + parallel_tokens_per_expert = parallel_tokens_per_expert.sum( + dim=0, dtype=torch.int + ) + parallel_bins = ops.inclusive_cumsum(parallel_tokens_per_expert, 0) + parallel_bins = ( + parallel_bins.view(1) if not len(parallel_bins.size()) else parallel_bins + ) + + # Calculate expert capacity + expert_capacity = expert_capacity_fn( + tokens_received, + top_k, + experts_per_rank_val, + expert_parallel_group, + moe_capacity_factor, + moe_expert_model_parallelism, + ) + if expert_capacity == 0: + expert_capacity = torch.max(parallel_tokens_per_expert).item() + + # Locally permute the tokens and perform the expert computation. + # Block to make sure that the cross-device permutation is complete. + if mlp_impl == "grouped": + # GroupedMLP requires counts on CPU. We can use the tensor already + # moved to CPU for the prior all_to_all, which avoids an extra + # device synchronization. + parallel_tokens_per_expert = parallel_tokens_per_expert_cpu.sum( + dim=0, + dtype=torch.int, + ) + + # Step 5: Expert computation + parallel_x_handle.wait() + + parallel_x = permute_and_compute( + parallel_x, + parallel_tokens_per_expert, + parallel_indices, + parallel_bin_ids, + None, # expert_weights + parallel_bins, + expert_capacity, + top_k=1, + w1=w1, + w2=w2, + w1_bias=w1_bias, + w2_bias=w2_bias, + gradient_scale=gradient_scale, + alpha=alpha, + ) + + # Step 6: Reverse communication - send results back + x, _ = _layers.all_to_all.all_to_all( + parallel_x, send_counts, recv_counts, expert_parallel_group + ) + + # Step 7: Reduce across hidden sharding dimension + shape = (hidden_sharding_deg, -1, hidden_size) + x = x.view(shape).sum(dim=0) + + # Step 8: Final local unpermutation + x = ops.scatter(x, indices, bin_ids, expert_weights, bins, top_k) + + return x, tokens_per_expert.flatten() + + +def moe_forward( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: Optional[torch.Tensor], + moe_top_k: int, + moe_num_experts: int, + moe_jitter_eps: float = None, + moe_normalize_expert_weights: int = None, + uniform_expert_assignment: bool = False, + training: bool = False, + w1: torch.Tensor = None, + w2: torch.Tensor = None, + w1_bias: torch.Tensor = None, + w2_bias: torch.Tensor = None, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + sort_end_bit: int = 0, + expert_parallel_group: torch.distributed.ProcessGroup = None, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = False, + forward_fn: Any = None, + hidden_size: int = None, + mlp_impl: str = "grouped", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + # Route tokens to experts + logits, expert_weights, expert_indices = route_tokens( + x, + router_weight, + router_bias, + moe_top_k, + moe_num_experts, + moe_jitter_eps, + moe_normalize_expert_weights, + uniform_expert_assignment, + training, + ) + + # Create router scores for output + router_scores = ( + torch.zeros_like(logits) + .scatter_(1, expert_indices, expert_weights) + .transpose(0, 1) + ) + + in_shape = x.size() + + # Prepare forward function arguments + forward_args = { + "x": x, + "expert_weights": expert_weights, + "top_experts": expert_indices, + "w1": w1, + "w2": w2, + "w1_bias": w1_bias, + "w2_bias": w2_bias, + "gradient_scale": gradient_scale, + "alpha": alpha, + "sort_end_bit": sort_end_bit, + "top_k": moe_top_k, + "num_experts": moe_num_experts, + "expert_parallel_group": expert_parallel_group, + "moe_capacity_factor": moe_capacity_factor, + "moe_expert_model_parallelism": moe_expert_model_parallelism, + "mlp_impl": mlp_impl, + } + + # Add hidden_size for parallel forward + if moe_expert_model_parallelism and hidden_size is not None: + forward_args["hidden_size"] = hidden_size + elif moe_expert_model_parallelism and hidden_size is None: + # Infer hidden_size from input shape + forward_args["hidden_size"] = x.shape[-1] + + # Compute expert outputs + x, tokens_per_expert = forward_fn(**forward_args) + + # Save load balancing loss if needed + moe_loss_weight = 0.0 # Can be made configurable + if training and moe_loss_weight > 0: + save_load_balancing_loss((tokens_per_expert, logits)) + + # Restore original shape + x = x.view(in_shape) + + return x, expert_weights, router_scores + + +def moe_forward_with_shared_expert( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: Optional[torch.Tensor], + moe_top_k: int, + moe_num_experts: int, + moe_jitter_eps: float = None, + moe_normalize_expert_weights: int = None, + uniform_expert_assignment: bool = False, + training: bool = False, + w1: torch.Tensor = None, + w2: torch.Tensor = None, + w1_bias: torch.Tensor = None, + w2_bias: torch.Tensor = None, + gradient_scale: Optional[float] = None, + alpha: float = 1.702, + sort_end_bit: int = 0, + expert_parallel_group: torch.distributed.ProcessGroup = None, + moe_capacity_factor: float = 1.0, + moe_expert_model_parallelism: bool = False, + forward_fn: Any = None, + hidden_size: int = None, + mlp_impl: str = "grouped", + # Shared expert parameters + shared_up_proj_weight: Optional[torch.Tensor] = None, + shared_down_proj_weight: Optional[torch.Tensor] = None, + shared_up_proj_bias: Optional[torch.Tensor] = None, + shared_down_proj_bias: Optional[torch.Tensor] = None, + shared_expert_weighted_sum: bool = False, + shared_activation_fn: Optional[Any] = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + # First, compute regular MoE forward pass + expert_out, expert_weights, router_scores = moe_forward( + x=x, + router_weight=router_weight, + router_bias=router_bias, + moe_top_k=moe_top_k, + moe_num_experts=moe_num_experts, + moe_jitter_eps=moe_jitter_eps, + moe_normalize_expert_weights=moe_normalize_expert_weights, + uniform_expert_assignment=uniform_expert_assignment, + training=training, + w1=w1, + w2=w2, + w1_bias=w1_bias, + w2_bias=w2_bias, + gradient_scale=gradient_scale, + alpha=alpha, + sort_end_bit=sort_end_bit, + expert_parallel_group=expert_parallel_group, + moe_capacity_factor=moe_capacity_factor, + moe_expert_model_parallelism=moe_expert_model_parallelism, + forward_fn=forward_fn, + hidden_size=hidden_size, + mlp_impl=mlp_impl, + ) + + # If shared expert weights provided, compute shared expert output + if shared_up_proj_weight is not None and shared_down_proj_weight is not None: + shared_expert_out = shared_mlp_forward( + x=x, + up_proj_weight=shared_up_proj_weight, + down_proj_weight=shared_down_proj_weight, + up_proj_bias=shared_up_proj_bias, + down_proj_bias=shared_down_proj_bias, + activation_fn=shared_activation_fn, + gradient_scale=gradient_scale, + ) + + # Combine expert outputs + combined_out = combine_expert_shared_outputs( + shared_expert_out=shared_expert_out, + expert_out=expert_out, + shared_expert_weighted_sum=shared_expert_weighted_sum, + moe_top_k=moe_top_k, + ) + + return combined_out, expert_weights, router_scores + + # Return regular MoE output if no shared expert + return expert_out, expert_weights, router_scores + + +def create_shared_expert_weights( + hidden_size: int, + shared_expert_hidden_size: int, + device: torch.device, + dtype: torch.dtype, + init_method: Any, + output_layer_init_method: Any = None, +) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + + if output_layer_init_method is None: + output_layer_init_method = init_method + + # Create weight tensors + up_proj_weight = torch.empty( + shared_expert_hidden_size, + hidden_size, + device=device, + dtype=dtype, + ) + down_proj_weight = torch.empty( + hidden_size, + shared_expert_hidden_size, + device=device, + dtype=dtype, + ) + + # Initialize weights + init_method(up_proj_weight) + output_layer_init_method(down_proj_weight) + + # No bias by default + return up_proj_weight, down_proj_weight, None, None + + +# HACK: Extract device_mesh from pre-hook closure - required for transformers integration +# This exists because device_mesh is trapped in hook closures with no model attribute +# Fragile - breaks if hook structure changes or Python internals change +# TODO: Replace with a more robust solution when available +def get_device_mesh(model): + # Extract device_mesh from child's unused pre_hook closure + try: + # Find the pre-hook that contains 'device_mesh' in its closure + hook = next( + h + for h in model.experts._forward_pre_hooks.values() + if "device_mesh" in h.__code__.co_freevars + ) + # Extract the device_mesh from the closure + return hook.__closure__[ + hook.__code__.co_freevars.index("device_mesh") + ].cell_contents + except Exception: + return None + + +class MegaBlocksMoeMLP(torch.nn.Module): + can_torch_compile: bool = True + + def forward(self, x: torch.Tensor) -> torch.Tensor: + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + gradient_scale = getattr(self.experts, "gradient_scale", None) + alpha = getattr(self.experts, "alpha", 1.0) + moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0) + moe_jitter_eps = getattr(self.experts, "jitter_eps", None) + moe_normalize_expert_weights = getattr( + self.experts, "normalize_expert_weights", None + ) + uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False) + + expert_parallel_group = getattr(self, "expert_parallel_group", None) + if expert_parallel_group is None: + device_mesh = get_device_mesh(self) + expert_parallel_group = device_mesh.get_group() if device_mesh else None + + has_parallel = ( + expert_parallel_group is not None + and dist.is_initialized() + and dist.get_world_size(expert_parallel_group) > 1 + ) + forward_fn = parallel_forward_once if has_parallel else forward_once + + sort_end_bit = max( + int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1 + ) + mlp_impl = getattr(self, "mlp_impl", "grouped") + output, expert_weights_out, *_ = moe_forward( + x=x, + router_weight=self.router.weight, + router_bias=self.router.bias, + moe_top_k=moe_top_k, + moe_num_experts=moe_num_experts, + moe_jitter_eps=moe_jitter_eps, + moe_normalize_expert_weights=moe_normalize_expert_weights, + uniform_expert_assignment=uniform_expert_assignment, + training=self.training, + w1=self.experts.gate_up_proj, + w2=self.experts.down_proj, + w1_bias=self.experts.gate_up_proj_bias, + w2_bias=self.experts.down_proj_bias, + gradient_scale=gradient_scale, + alpha=alpha, + sort_end_bit=sort_end_bit, + expert_parallel_group=expert_parallel_group, + moe_capacity_factor=moe_capacity_factor, + moe_expert_model_parallelism=has_parallel, + forward_fn=forward_fn, + hidden_size=self.experts.hidden_size, + mlp_impl=mlp_impl, + ) + return output, expert_weights_out + + +# Export main classes +__all__ = ["MegaBlocksMoeMLP", "MegaBlocksMoeMLPWithSharedExpert"] + + +class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP): + + def __init__(self): + super().__init__() + # Shared expert weights will be set by the user + self.shared_up_proj_weight = None + self.shared_down_proj_weight = None + self.shared_up_proj_bias = None + self.shared_down_proj_bias = None + self.shared_expert_weighted_sum = False + self.shared_activation_fn = None + + def set_shared_expert_weights( + self, + up_proj_weight: torch.Tensor, + down_proj_weight: torch.Tensor, + up_proj_bias: Optional[torch.Tensor] = None, + down_proj_bias: Optional[torch.Tensor] = None, + weighted_sum: bool = False, + activation_fn: Optional[Any] = None, + ): + self.shared_up_proj_weight = up_proj_weight + self.shared_down_proj_weight = down_proj_weight + self.shared_up_proj_bias = up_proj_bias + self.shared_down_proj_bias = down_proj_bias + self.shared_expert_weighted_sum = weighted_sum + self.shared_activation_fn = activation_fn + + def forward(self, x: torch.Tensor) -> torch.Tensor: + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + gradient_scale = getattr(self.experts, "gradient_scale", None) + alpha = getattr(self.experts, "alpha", 1.0) + moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0) + moe_jitter_eps = getattr(self.experts, "jitter_eps", None) + moe_normalize_expert_weights = getattr( + self.experts, "normalize_expert_weights", None + ) + uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False) + + expert_parallel_group = getattr(self, "expert_parallel_group", None) + if expert_parallel_group is None: + device_mesh = get_device_mesh(self) + expert_parallel_group = device_mesh.get_group() if device_mesh else None + + has_parallel = ( + expert_parallel_group is not None + and dist.is_initialized() + and dist.get_world_size(expert_parallel_group) > 1 + ) + forward_fn = parallel_forward_once if has_parallel else forward_once + + sort_end_bit = max( + int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1 + ) + mlp_impl = getattr(self, "mlp_impl", "grouped") + + output, expert_weights_out, *_ = moe_forward_with_shared_expert( + x=x, + router_weight=self.router.weight, + router_bias=self.router.bias, + moe_top_k=moe_top_k, + moe_num_experts=moe_num_experts, + moe_jitter_eps=moe_jitter_eps, + moe_normalize_expert_weights=moe_normalize_expert_weights, + uniform_expert_assignment=uniform_expert_assignment, + training=self.training, + w1=self.experts.gate_up_proj, + w2=self.experts.down_proj, + w1_bias=self.experts.gate_up_proj_bias, + w2_bias=self.experts.down_proj_bias, + gradient_scale=gradient_scale, + alpha=alpha, + sort_end_bit=sort_end_bit, + expert_parallel_group=expert_parallel_group, + moe_capacity_factor=moe_capacity_factor, + moe_expert_model_parallelism=has_parallel, + forward_fn=forward_fn, + hidden_size=self.experts.hidden_size, + mlp_impl=mlp_impl, + # Shared expert parameters + shared_up_proj_weight=self.shared_up_proj_weight, + shared_down_proj_weight=self.shared_down_proj_weight, + shared_up_proj_bias=self.shared_up_proj_bias, + shared_down_proj_bias=self.shared_down_proj_bias, + shared_expert_weighted_sum=self.shared_expert_weighted_sum, + shared_activation_fn=self.shared_activation_fn, + ) + return output, expert_weights_out + + +# Patch for XPU or CPU support +if hasattr(torch, "xpu") and torch.xpu.is_available(): + from .xpu_fused_moe import MegaBlocksMoeMLP + +from .cpu_moe_cpp import CPUMegaBlocksMoeMLP diff --git a/build/torch214-cxx11-cu132-x86_64-linux/metadata.json b/build/torch214-cxx11-cu132-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..89f06447d7124fcd885ff71bd5693a0fecf236de --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/metadata.json @@ -0,0 +1,120 @@ +{ + "name": "megablocks", + "id": "_megablocks_cuda_addf474", + "version": 1, + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.3", + "11.0", + "12.0", + "12.1+PTX", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "o7Y9XLXezcJogNTNoC4HMHsd2ZZ3CTwzYftkUthOtno=", + "_grouped_gemm_triton/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "_grouped_gemm_triton/_triton_kernels/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "_grouped_gemm_triton/_triton_kernels/gmm.py": "11SvNqMt/LDCgjJBOmZlHg9yn9NbGw1UYKzNR/JKtcQ=", + "_grouped_gemm_triton/adapter.py": "tP00AypBJn4rDqrBDMMwBb6MwJ6s9dg18GNQxpLiVoA=", + "_grouped_gemm_triton/configs.py": "aHXlKELHIy8J82m1kClQRxDkad4LCL8BOeEpq4/IqKQ=", + "_grouped_gemm_triton/gmm.py": "1iqK7N3vsahgQ4bbHsOMA8hoR5uQhh1PmbFWjOLSKDo=", + "_grouped_gemm_triton/utils/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "_grouped_gemm_triton/utils/_triton/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "_grouped_gemm_triton/utils/_triton/arch_info.py": "Vo7tziqn9GLA8S4xrKT2555aKFxo75Oja4ek2zNReuY=", + "_grouped_gemm_triton/utils/_triton/pid_preprocessing.py": "Lx1S7li2xArbeql/pEFfyul7yU61agwuYHTsZ7jG0os=", + "_grouped_gemm_triton/utils/gmm_common.py": "pQMeVpP9kk/bixwJE309iceJPa8H7mZV8b6+OtbQvl8=", + "_grouped_gemm_triton/utils/logger.py": "CXuJZCBzQDTrW80OyJ2IuhbHhayLy12Ake1KGxiuYIA=", + "_layers/__init__.py": "bwEDLk8kbLuT1BWZPPDVUiDsWsJ3wC3nI+O+dSGAoSU=", + "_layers/activation_fn.py": "lJrZcw+1oLw0XOyjYb6AfmxX1n6g2Snbph6H3p5D0oU=", + "_layers/all_to_all.py": "SZMHtoTy6B76caEE0shoK1y4iHqWp+lWqdR/30d9oBQ=", + "_layers/arguments.py": "wYFjDwBidH5MjMr4cZGDEcRvW1IgdLQvcYt/qt5S0dk=", + "_layers/common.py": "8m+bNabk+jB+C8o797GdNpGzx8m3nbnG66FQdnDZRVI=", + "_layers/dmlp_registry.py": "z7iqmxiJx9eACdaHOCmnY9EGKV/zePhqcf8FVboaoK8=", + "_layers/dmoe.py": "bOaQ0PFOOLu45HxJNRFuKPB8VedUzsduekkqJV6WrQ8=", + "_layers/gelu.py": "fcq4OBCSCcPjCk7PYZydIpbdn/iUyMtL2TqhonHqt7k=", + "_layers/glu.py": "D27/lTdF05N6TdSPoyW1MWgTGyaXJlUCo5uos54zabU=", + "_layers/memory_test.py": "4jm72yP6WBV5zelWQGJQIHHZLPisFEtar+moOSnxXMk=", + "_layers/mlp.py": "xx51A1pGeBt/B8m7pI5hjhVqhEllpxvldulXYzFmA3Y=", + "_layers/moe.py": "vVqEx4Ppgw+QbLHOJo5uSI/qXDikA99ZGApPOcBJtmg=", + "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=", + "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=", + "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=", + "_megablocks_cuda_addf474.abi3.so": "PIfCfE194zjfsaPspPY/uFhbovFERYw70ZZuRhbi5GI=", + "_ops.py": "WXlV4XwtKQjdQ6nap3k9TalaHkK0ypeSTshOrV6Hczo=", + "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=", + "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=", + "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=", + "benchmark_util.py": "WFsGWGaJvWvwdx4w33WCfNzSJkVuQrZmZCaAodJsws8=", + "cpu_fused_moe.py": "XN2eQBFaihbJrZ/cu9jS8h5dcoiUNh82gjJQUARcBYA=", + "cpu_moe_cpp.py": "dvv4Ynn/FWObs421L4hAdYZPiPs2BS/sbEOuMMIUPVM=", + "grouped_gemm/__init__.py": "fxW9amtzy/Z0V/OXSYMmibD7XvpYBGEv81KMxvocFKY=", + "grouped_gemm/backend.py": "Yd//ihpvu/pxEoGBq3iD1zMggnYv6quwZMdg1R4kUEI=", + "grouped_gemm/ops.py": "Kj7yRgv7afB5SVwXIkN9RNIF2kxD2qrZ0dVsKX5ph2s=", + "grouped_gemm_util.py": "bVPrdsEtBVJJ+v/bNrGNf/QHNlVL4VxHg92HvaL3oq4=", + "layers.py": "mNTlX2oKJyHBJllRd24lPNdGLa2H9cEUgVLvkZJpoNo=", + "ops/__init__.py": "MHQsiMEgkKeZCa8TAYFMGDFnwFIgijfAN7DhMUt/48U=", + "ops/all_to_all_benchmark.py": "sxXSCSDdsh+62IWpKNTyBphLTngY1jlgJ1WgNooSEO8=", + "ops/binned_gather.py": "pkWne+ThfiTR/bMbDs2QDimKxAI4klqD1FdvhytE7lU=", + "ops/binned_scatter.py": "5NxocRmgYDkTAntDtXBlF4c2r8rfbLVnBRPA6f4/L4U=", + "ops/cumsum.py": "upuWnnp125kQhJrIMLtG09cI1zzxEy06P5wYhOwbpI4=", + "ops/gather.py": "a8j/btP1VCdJ7F7FC58MFPxFuH9tUhjMuQnEzwjAgGw=", + "ops/histogram.py": "v5+pHF+/fjKQAkNLIsHyljMUakMTdW5GzA094PgRoGk=", + "ops/histogram_benchmark.py": "oX8VolN+/PVCWaPHkY0k7SNn4QpadHi80Sla/3NScNI=", + "ops/matmul_benchmark.py": "SZ3L/CatnT4J5PyOU/Ii/b5y5fw3HNEaR/xY6PjrJ8A=", + "ops/padded_gather.py": "VtPH3Ge1jodeug5MWwXnxnmGIi0D1MOOaXi2CcU8lpI=", + "ops/padded_scatter.py": "UO1BwsaezK3AikD8CHoJOEjmC4h+H5fcOX/bvO8FWXw=", + "ops/padded_scatter_benchmark.py": "w34qlIc3Nn0AgB2k7iIFqQXgDNplMVqCr5i8GrgrvZ0=", + "ops/permute_benchmark.py": "ISYPkVZknTnfB8fhUVloSyfEOe7ApKQKtNPvEq3P9qg=", + "ops/repeat.py": "j9l1A/4LVhcFuSAaneFlImq7RGNXie0zfu2nMlEhHI4=", + "ops/replicate.py": "NN5843jJz0OdfYS6YaPCpRD8dGk66S3SHTzR51PrwUI=", + "ops/round_up.py": "/CQSCYmjW0Mrv8TFdfZs/ZpXHHz18b3E4aLHvmfUxOY=", + "ops/scatter.py": "YgM76I/dYAzRNN+jzAdaRWan56GiwuCZRR7xy3RAY8Y=", + "ops/sort.py": "35D3oG4AT0Jd6+awwfoL+GmkcNtdRLno7LcSFidqKts=", + "ops/sort_benchmark.py": "Bo985Xsxq/F0PMU92xKCo+srSKVdheQLey2GqFlwZXY=", + "ops/stk_autocast.py": "SQArC9BFDXVManccbEWibSg0E+/R3NQlkp0cB4mIjXo=", + "ops/sum.py": "fcjWM5Aiz1wQ+UYBN2Y/OFILsiGKSLgmlkeqfVRccD4=", + "ops/topology.py": "+lfjfxDjB+st3K8wZnRY5fEanSIWT+3SJncx9bu+6P0=", + "stk/__init__.py": "QFugfJQ5q1Wdn/+VQyfs1YKns79i3FiywiLevg0ymXo=", + "stk/backend/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=", + "stk/backend/autocast.py": "ppk7mL/WQUHaLqEgsW3h5ZGhldC8WyIKmaIS6kAe4lA=", + "stk/backend/sputnik.py": "Mpbo0M1ES39PdaIo65ef3/mQ8XHtLbO3ayqufaIZLiw=", + "stk/backend/triton_kernels.py": "RCthwX2feppkqBK3pvm+/37nhF89lLHgP12c6xrfxXU=", + "stk/matrix.py": "Xxo1RwOtKFgVXExhlAoIC67V+eWKZRClDdbBfzUgfy0=", + "stk/ops/__init__.py": "Vqbt6C0qWEPaWypsCoaHGoZRbdPcFrLg1ahlAlp4l+Y=", + "stk/ops/eltwise_ops.py": "O9TqOMvZ2NZ7KNFaeuCmJOfg1Hz5ad8rvYwbGAH7/+0=", + "stk/ops/eltwise_ops_test.py": "rD1oPPu+vXHlyJnEJOpM9L5wdCeiDbn98CbCH8mQ0wE=", + "stk/ops/linear_ops.py": "ZWOpLTXMdw/ZTyxL4W74CSSZQaUKvtfiz3ev0OnaVQ0=", + "stk/ops/linear_ops_test.py": "kl8nMX1EAvZ6lF7oN4U4bW2z0hera0YxqTPiNljLmzk=", + "stk/ops/matrix_ops.py": "+hhUk7IZZFVVzIk1pttWFFSOtKwlGEBVR2V4eyL27K4=", + "stk/ops/matrix_ops_test.py": "Bb4fk4+IzNC0zOFq0Tz/JwNaLWyhtWEDgEUzyVdbnrE=", + "stk/random/__init__.py": "Gr/+Gl20jJ6fJVRfoAUxEcCPeChILQzrzhKVVAWQsRI=", + "stk/random/random_ops.py": "AluIDj3Vo2ReKAsZz07yydbtndDbJoglekRGIaCoRZE=", + "stk/random/random_ops_test.py": "B+a4Zb+OPhBtpVA48oc09sBL/AAQO9j/YUFI0nnXxzk=", + "xpu_fused_moe.py": "sPmuRSpkH1Bmk4hdpz53us8Qpc6yinT9rkld3T8WreI=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "commit": "7d2828e6592de641b44fb8eb896719101a2ab2fb", + "dirty": false + }, + "kernel": { + "commit": "addf4741fafb55b4c101b15985dfaeea7c7c2f3a", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch214-cxx11-cu132-x86_64-linux/metadata.json.sigstore b/build/torch214-cxx11-cu132-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..59dbf3b1fde2d1368fbcc94abdd443ae0f953645 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtCgAwIBAgIUfzH2ofHDeZXI81YX4xJuzX43qQ4wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI3MDcxNDEwWhcNMjYwODI3MDcyNDEwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEaq/e1nxnv4sO6b/61sWRlaNwsWfxYifdxGjTGfVm2b4F/GgqinrT30D8U3OovK4kXdz6E0C/ildLdEf5HM6qo6OCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUq3kdpctylI77GPbq7pSJ20XzWk4wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGFkZGY0NzQxZmFmYjU1YjRjMTAxYjE1OTg1ZGZhZWVhN2M3YzJmM2EwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzMwNDUzOTU5NTIvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoEIRr8oAAAQDAEYwRAIgF06z0URWMd1rUxpOfvNRe56x9lszjy6chJFoMgmCefUCIEMFdQzxUY3d7JhrFF5QIL+8s6GefQUQguMA156zgGLcMAoGCCqGSM49BAMDA2gAMGUCMBXOeCA9jEkgptX0uo5vLVVKt/Qet89Jb5BV4T4GSod+PomLMXMektStFqKQDAIvhwIxAIVmu4LVMZYXWBEjM75p4OYkFFuAr97pXXELY9Zhe4nGJolOKhEBHfS4gPKPi54H4Q=="}, "tlogEntries":[{"logIndex":"2614048124", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787814850", "inclusionPromise":{"signedEntryTimestamp":"MEQCIEqx6U66wkoI8TL5bDi9Z8W80zZKbcQcQLdXcLKgl2SSAiBKQ3c5QIZDAAlzJ1diEqsAwViswR6c2w78p13MWl+6pw=="}, "inclusionProof":{"logIndex":"2492143862", "rootHash":"uZarBYAFWByVwpjjRYjhPJNMNOnjO4IpZtq33cFi8yM=", "treeSize":"2492143865", "hashes":["GPm3nE03g20EZnpZWJGyg+rsKcXPoXynfIlWAqfOg9o=", "wUwHSZKDtFbHA8Rzjw0l+PcXKSpX2W9+rBzHxLRSfPc=", "QUIlvNy71+1D03k0/nyvDDAOJvH+ja7ywpjP7K+WIFI=", "whajdUkmjChxHWlKSx1m/JfxqBJ/hWtGTUGDj5thp6c=", "s/ul8fkEZwo5w+ae94EyzhvUuJL8iwOsXeyI0PToxRI=", "bDSsLiNYD4/CqLxGPWCqjygm3yqW4u8JWyaHBGv3Qhc=", "3AU20apjKw6TITqX7VoCm91fqdzXCdxRpW6Ioz4adHs=", "ZbjQmD+8kOp7X+/HUJRaLiRq7bXipiSM2D1iUO+dwzg=", "VFlTxnDU4W+k/WA3ARTvzCiqVt1yxQIMUyUzMCugJ2Q=", "CFNMzmJ8N+uNx37P2asMqvrW9tw4axAVHnXLRO39tFs=", "NypqAbB+mNix4l1ahGeP9fRWiQzZCm4P1664makjjq4=", "8SoaeA7IeJPoIiveQbBfHCWXSDQabIk2PW+FT49NSEQ=", "VGrFt5JvlIUQvoGEdBBxrCe3lYbu581OhgLq/Xdyat4=", "7TfDCaMWgEwSJEZd/e33b24919rF4IDJLmbqNNRNhZs=", "0HfdxDt/zGugZLIuHasrdLEW9s8OaTykSSqBEDFQXdw=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2492143865\nuZarBYAFWByVwpjjRYjhPJNMNOnjO4IpZtq33cFi8yM=\n\n— rekor.sigstore.dev wNI9ajBGAiEAwIXkZlziP0dgAw2vljjDN11x7N+ijbq/xmZ39lLs46QCIQCSJ+1PaRreXmb6HUOfzUQDguREIHCdxWtDt47GZ9bdpA==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJlZjFlYTY3YTFlM2E0NDc2MTIyN2JmMmZkZWMwN2QzZmFmZmM4N2QwMWRhYjNiOTI0ZDAxMzU4MGMzZTQ0YzljIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJQ00vNVE5d0hSbVpsVkV0ZCs1VG02ODhvVXJ2RkRZcmExTDZhTUpIcEcvUUFpQXlDK0RxLzZBd1IyT1JTNWtRWG1VOVprTHVzRlRHVEl3bDgyazBGNkt6TkE9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SRFowRjNTVUpCWjBsVlpucElNbTltU0VSbFdsaEpPREZaV0RSNFNuVjZXRFF6Y1ZFMGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1ROTlJHTjRUa1JGZDFkb1kwNU5hbGwzVDBSSk0wMUVZM2xPUkVWM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZoY1M5bE1XNTRiblkwYzA4MllpODJNWE5YVW14aFRuZHpWMlo0V1dsbVpIaEhhbFFLUjJaV2JUSmlORVl2UjJkeGFXNXlWRE13UkRoVk0wOXZka3MwYTFoa2VqWkZNRU12YVd4a1RHUkZaalZJVFRaeGJ6WlBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZ4TTJ0a0NuQmpkSGxzU1RjM1IxQmljVGR3VTBveU1GaDZWMnMwZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVpWMUpyV21wUk0wNUVSbTFaVjFwcFRsUldhVTVIVFhoTlJFWnBDazFVVlRWUFJGWnJXbTFHYkZwWFJUTlplbVJxVFcxWmVsbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWxOVkZVMVQwUldhMXB0Um13S1dsZEZNMWw2WkdwTmJWbDZXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWtLVFZSVk5VOUVWbXRhYlVac1dsZEZNMWw2WkdwTmJWbDZXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZEdhMXBIV1RBS1RucFJlRnB0Um0xWmFsVXhXV3BTYWsxVVFYaFpha1V4VDFSbk1WcEhXbWhhVjFab1RqSk5NMWw2U20xTk1rVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWsxM1RrUlZlazlVVlRWT1ZFbDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUZTVkp5T0c5QlFVRlJSQXBCUlZsM1VrRkpaMFl3Tm5vd1ZWSlhUV1F4Y2xWNGNFOW1kazVTWlRVMmVEbHNjM3BxZVRaamFFcEdiMDFuYlVObFpsVkRTVVZOUm1SUmVuaFZXVE5rQ2pkS2FISkdSalZSU1V3ck9ITTJSMlZtVVZWUlozVk5RVEUxTm5wblIweGpUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlaMEZOUjFWRFRVSllUMlZEUVRrS2FrVnJaM0IwV0RCMWJ6VjJURlpXUzNRdlVXVjBPRGxLWWpWQ1ZqUlVORWRUYjJRclVHOXRURTFZVFdWcmRGTjBSbkZMVVVSQlNYWm9kMGw0UVVsV2JRcDFORXhXVFZwWldGZENSV3BOTnpWd05FOVphMFpHZFVGeU9UZHdXRmhGVEZrNVdtaGxORzVIU205c1QwdG9SVUpJWmxNMFoxQkxVR2sxTkVnMFVUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgff2Uj5C8HGoD54bgRRO79TVRBe1xt641Kb/STBTbDDYCFQDewwfhR3gl70Emguei9atUbjHoNhgPMjAyNjA4MjcwNzE0MTBaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgyNzA3MTQxMFowLwYJKoZIhvcNAQkEMSIEILJkBQC7owW+M/S/TqulUIEV7FJbo1NgSWCVYO1RRKxvMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQD+MrvoOoguDDrBoNu6j9OEPQ+HNp3SjliPbcRFubTSam+lkg+qf9oBIipXstDX/c8CME/nCcidtDVC0ytcnty3ZiXtavEOznih2qaadn02oqIeq33xSzDVnqAAaEZsNvubig=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"7x6meh46RHYSJ78v3sB9P6/8h9AdqzuSTQE1gMPkTJw="}, "signature":"MEQCICM/5Q9wHRmZlVEtd+5Tm688oUrvFDYra1L6aMJHpG/QAiAyC+Dq/6AwR2ORS5kQXmU9ZkLusFTGTIwl82k0F6KzNA=="}} \ No newline at end of file diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b944080df810d0b0cfc571f3009b0098a651f9b7 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/__init__.py @@ -0,0 +1,35 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from .binned_gather import binned_gather +from .binned_scatter import binned_scatter +from .cumsum import exclusive_cumsum, inclusive_cumsum +from .gather import gather +from .histogram import histogram +from .padded_gather import padded_gather +from .padded_scatter import padded_scatter +from .repeat import repeat +from .replicate import replicate +from .round_up import round_up +from .scatter import scatter +from .sort import sort +from .sum import sum +from .topology import topology + +__all__ = [ + 'binned_gather', + 'binned_scatter', + 'exclusive_cumsum', + 'inclusive_cumsum', + 'gather', + 'histogram', + 'padded_gather', + 'padded_scatter', + 'repeat', + 'replicate', + 'round_up', + 'scatter', + 'sort', + 'sum', + 'topology', +] diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/all_to_all_benchmark.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/all_to_all_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..4c939818edca3345f6344bbc7cef07ffe3cd0181 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/all_to_all_benchmark.py @@ -0,0 +1,63 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.distributed as dist + +# from megablocks import benchmark_util +# from megablocks.layers.all_to_all import all_to_all + +from .. import benchmark_util +from .._layers.all_to_all import all_to_all + +_ALL_TO_ALL_BENCHMARK = ( + (8, 1024), + (16, 1024), + (32, 1024), + (64, 1024), + (128, 1024), + (256, 1024), + (512, 1024), + (1024, 1024), + (2 * 1024, 1024), + (4 * 1024, 1024), + (8 * 1024, 1024), + (16 * 1024, 1024), + (32 * 1024, 1024), + (64 * 1024, 1024), + (128 * 1024, 1024), + (256 * 1024, 1024), + (512 * 1024, 1024), + (1024 * 1024, 1024), +) + + +def benchmark_all_to_all(group, sl, hs): + world_size = dist.get_world_size(group) + assert (sl % world_size) == 0 + send_recv_sizes = [sl // world_size] * world_size + + x = torch.randn((sl, hs)).cuda().half() + + details = { + 'world_size': world_size, + 'message_size (B)': send_recv_sizes[0] * hs * 2, # 2B elements. + } + + def benchmark(): + return all_to_all(x, send_recv_sizes, send_recv_sizes, group) + + time, std = benchmark_util.benchmark_function(benchmark) + + if dist.get_rank(group) == 0: + benchmark_util.log_benchmark('All-To-All', details, time, std) + + +if __name__ == '__main__': + assert dist.is_available() + group = dist.init_process_group(backend='nccl') + local_rank = dist.get_rank(group) + torch.cuda.set_device(local_rank) + + for args in _ALL_TO_ALL_BENCHMARK: + benchmark_all_to_all(group, *args) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/binned_gather.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/binned_gather.py new file mode 100644 index 0000000000000000000000000000000000000000..189a7fa3518d660f29ea32e7a04827164af98d60 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/binned_gather.py @@ -0,0 +1,37 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for binned_gather kernel. +class BinnedGatherOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bins: torch.Tensor, + bin_size: int, + top_k: int, + ): + ctx.save_for_backward(indices, bins) + ctx.top_k = top_k + return kernels.binned_gather(x, indices, None, bins, bin_size, top_k) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + indices, bins = ctx.saved_tensors + out = kernels.binned_scatter(grad, indices, None, bins, ctx.top_k) + return out, None, None, None, None + + +binned_gather = BinnedGatherOp.apply diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/binned_scatter.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/binned_scatter.py new file mode 100644 index 0000000000000000000000000000000000000000..cb937c0c106662ce8108c1cb926f8f063b163d3d --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/binned_scatter.py @@ -0,0 +1,59 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for binned_scatter kernel. +class BinnedScatterOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + top_k: int, + ): + assert len(x.size()) == 3 + ctx.bin_size = x.size(1) + ctx.top_k = top_k + + # TODO(tgale): Don't save 'x' for backwards if we don't need to + # calculate the gradient w.r.t. 'weights'. + ctx.save_for_backward(x, indices, weights, bins) + return kernels.binned_scatter(x, indices, weights, bins, top_k) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + x, indices, weights, bins = ctx.saved_tensors + out = kernels.binned_gather( + grad, + indices, + weights, + bins, + ctx.bin_size, + ctx.top_k, + ) + + wgrad = None + if ctx.needs_input_grad[2]: + wgrad = kernels.binned_scatter_wgrad( + x, + grad, + indices, + bins, + ctx.top_k, + ) + return out, None, wgrad, None, None + + +binned_scatter = BinnedScatterOp.apply diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/cumsum.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/cumsum.py new file mode 100644 index 0000000000000000000000000000000000000000..e2b7572391e20045d335cf7337246e8a9b9f57ef --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/cumsum.py @@ -0,0 +1,52 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + # import megablocks_ops as ops # type: ignore + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + + +# Autograd wrappers for cumsum kernels. +# NOTE: Does not support gradients. +class ExclusiveCumsumOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, dim: int): + if len(x.size()) == 1: + x = x.view([1, -1]) + out = torch.empty_like(x) + ops.exclusive_cumsum(x, 1, out) + return out.squeeze() + out = torch.empty_like(x) + ops.exclusive_cumsum(x, dim, out) + return out + + +exclusive_cumsum = ExclusiveCumsumOp.apply + + +class InclusiveCumsumOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, dim: int) -> torch.Tensor: + if len(x.size()) == 1: + x = x.view([1, -1]) + out = torch.empty_like(x) + ops.inclusive_cumsum(x, 1, out) + return out.squeeze() + out = torch.empty_like(x) + ops.inclusive_cumsum(x, dim, out) + return out + + +inclusive_cumsum = InclusiveCumsumOp.apply diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/gather.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/gather.py new file mode 100644 index 0000000000000000000000000000000000000000..f1f87c1e7bed8d3589dd790805234976e0b05898 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/gather.py @@ -0,0 +1,38 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for gather kernel. +class GatherOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + bins: torch.Tensor, + top_k: int, + ): + ctx.save_for_backward(indices, bin_ids, bins) + ctx.top_k = top_k + return kernels.gather(x, indices, bin_ids, None, bins, top_k) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + + indices, bin_ids, bins = ctx.saved_tensors + out = kernels.scatter(grad, indices, bin_ids, None, bins, ctx.top_k) + return out, None, None, None, None, None + + +gather = GatherOp.apply diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/histogram.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/histogram.py new file mode 100644 index 0000000000000000000000000000000000000000..7b3f058ec373cbba7555704fb5e4212c3cc75d9d --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/histogram.py @@ -0,0 +1,27 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + + +# Autograd wrapper for histogram kernel. +# NOTE: Does not support gradients. +class HistogramOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, max_val: float): + return ops.histogram(x, max_val) + + +histogram = HistogramOp.apply diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/histogram_benchmark.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/histogram_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..a0b9c6047567b87a295979498142230d1b0c9411 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/histogram_benchmark.py @@ -0,0 +1,78 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import numpy as np +import torch +# from absl.testing import parameterized + +from .. import ops + +_HISTOGRAM_TESTS = ( + (16384, torch.int32, 2), + (16384, torch.int32, 4), + (16384, torch.int32, 8), + (16384, torch.int32, 16), + (16384, torch.int32, 32), + (16384, torch.int32, 64), + (16384, torch.int32, 128), + (16384, torch.int32, 256), +) + + +def benchmark_function(fn, iterations=10): + # Run once to get rid of startup overhead. + fn() + times = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + times = np.array(times) + return times.mean(), times.std(), times.max(), times.min() + + +def log_benchmark(arguments, mean_t, std_t): + print('=' * 60) + print('Benchmark Parameters:') + for (key, value) in arguments.items(): + print(f'{key} = {value}') + print('Results:') + print('mean / std = {:.2f}ms / {:.2f}ms'.format(mean_t, std_t)) + print('=' * 60) + + +# class HistogramBenchmark(parameterized.TestCase): +# +# @parameterized.parameters(*_HISTOGRAM_TESTS) +# def testHistogram(self, n, dtype, max_val): +# x = torch.randint(0, max_val, (n,)).cuda().to(dtype) +# +# mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.histogram(x, max_val),) +# arguments = { +# 'n': n, +# 'dtype': dtype, +# 'max_val': max_val, +# } +# log_benchmark(arguments, mean_t, std_t) +# +# @parameterized.parameters(*_HISTOGRAM_TESTS) +# def testTorchHistogram(self, n, dtype, max_val): +# x = torch.randint(0, 128, (n,)).cuda().to(dtype) +# +# mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.histc(x, max_val, 0, max_val - 1),) +# arguments = { +# 'n': n, +# 'dtype': dtype, +# 'max_val': max_val, +# } +# log_benchmark(arguments, mean_t, std_t) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/matmul_benchmark.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/matmul_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..4c4b9b8866ffed2eb769b77f2320c82e5491ae0e --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/matmul_benchmark.py @@ -0,0 +1,415 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + + +# import stk + +# try: +# import stk +# except ImportError: +# import warnings +# warnings.warn( +# 'Please add `stanford-stk` if megablocks/ops/matmul_benchmark.py is needed.', +# ) + +from .. import stk + +import torch +# from absl.testing import parameterized + +from .. import benchmark_util, ops + + +# Calling tensor.t() calls tensor.transpose(0, 1) which calls +# torch.as_strided(...). Circumvent this chain to avoid an overhead +# this adds. +def transpose_view(x): + return torch.as_strided( + x, + (x.shape[1], x.shape[0]), + (x.stride()[1], x.stride()[0]), + ) + + +_MATMUL_TESTS = ( + (64 * 1024, 512, 2048, 64), + (32 * 1024, 768, 3072, 64), + (8 * 1024, 1024, 4096, 64), + (4 * 2048, 4096, 4 * 4096, 4), +) + + +def log_benchmark(name, arguments, time, std, flops): + benchmark_util.log_benchmark(name, arguments, time, std) + print('flops = {:.2f}B'.format(flops / 1e9)) + print('throughput = {:.2f}T'.format(flops / 1e9 / time)) + print('=' * 60) + + +# class MatmulBenchmark(parameterized.TestCase): +# +# def build_sparse_matrix(self, x, padded_bins, fhs, ne): +# blocking = 128 +# padded_tokens, _ = x.size() +# assert padded_tokens % blocking == 0 +# assert fhs % blocking == 0 +# +# # Offsets for the sparse matrix. All rows have the +# # same number of nonzero blocks dictated by the +# # dimensionality of a single expert. +# block_rows = padded_tokens // blocking +# blocks_per_row = fhs // blocking +# offsets = torch.arange( +# 0, +# block_rows * blocks_per_row + 1, +# blocks_per_row, +# dtype=torch.int32, +# device=x.device, +# ) +# +# # Indices for the sparse matrix. The indices for +# # the intermediate matrix are dynamic depending +# # on the mapping of tokens to experts. +# column_indices = ops.topology( +# padded_bins, +# blocking, +# block_rows, +# blocks_per_row, +# ) +# data = torch.empty( +# column_indices.numel(), +# blocking, +# blocking, +# dtype=torch.float16, +# device=x.device, +# ) +# shape = (padded_tokens, fhs * ne) +# row_indices = stk.ops.row_indices(shape, data, offsets, column_indices) +# return stk.Matrix(shape, data, row_indices, column_indices, offsets) +# +# def build_input_matrix(self, sl, hs, ne): +# x = torch.randn((sl, hs)).cuda().half() +# +# # Assign tokens to experts uniformly. +# top_expert = torch.arange(0, sl).cuda().int() % ne +# +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(top_expert, ne) +# padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) +# padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# out = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, 1) +# return out, padded_bins +# +# def build_weight_matrix(self, ne, hs, fhs): +# return torch.randn((hs, ne * fhs)).cuda().half() +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_Fwd_SDD_NT(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# topo = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# w = transpose_view(w) +# +# def benchmark(): +# return stk.ops.sdd(x, w, topo) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0::Fwd::SDD::NT', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_GradX_DSD_NN(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# topo = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# +# def benchmark(): +# return stk.ops.dsd(topo, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0::GradX::DSD::NN', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_GradW_DSD_TN(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# topo = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# topo = topo.t() +# +# def benchmark(): +# return stk.ops.dsd(topo, x) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0::GradW::DSD::TN', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_Fwd_DSD_NN(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# x = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# +# def benchmark(): +# return stk.ops.dsd(x, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::Fwd::DSD::NN', +# arguments, +# mean_t, +# std_t, +# x.nnz * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_GradX_SDD_NT(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# x = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# out = stk.ops.dsd(x, w) +# w = transpose_view(w) +# +# def benchmark(): +# return stk.ops.sdd(out, w, x) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::GradX::SDD::NT', +# arguments, +# mean_t, +# std_t, +# x.nnz * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_GradW_DSD_TN(self, sl, hs, fhs, ne): +# x, padded_bins = self.build_input_matrix(sl, hs, ne) +# w = self.build_weight_matrix(ne, hs, fhs).t().contiguous() +# x = self.build_sparse_matrix(x, padded_bins, fhs, ne) +# out = stk.ops.dsd(x, w) +# x = x.t() +# +# def benchmark(): +# return stk.ops.dsd(x, out) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::GradW::DSD::TN', +# arguments, +# mean_t, +# std_t, +# x.nnz * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_Fwd_DDD_NT(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, hs)).cuda().half() +# w = torch.randn((ne, hs, fhs)).cuda().half() +# +# w = w.transpose(1, 2).contiguous() +# w = w.transpose(1, 2) +# +# def benchmark(): +# return torch.bmm(x, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0::Fwd:DDD::NT', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_GradX_DDD_NN(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, hs)).cuda().half() +# w = torch.randn((ne, hs, fhs)).cuda().half() +# out = torch.bmm(x, w) +# w = w.transpose(1, 2).contiguous() +# +# def benchmark(): +# return torch.bmm(out, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0:GradX:DDD::NN', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear0_GradW_DDD_TN(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, hs)).cuda().half() +# w = torch.randn((ne, hs, fhs)).cuda().half() +# out = torch.bmm(x, w) +# out = out.transpose(1, 2) +# +# def benchmark(): +# return torch.bmm(out, x) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '0:GradW:DDD::TN', +# arguments, +# mean_t, +# std_t, +# x.numel() * fhs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_Fwd_DDD_NN(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, fhs)).cuda().half() +# w = torch.randn((ne, fhs, hs)).cuda().half() +# +# def benchmark(): +# return torch.bmm(x, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::Fwd::DDD::NN', +# arguments, +# mean_t, +# std_t, +# x.numel() * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_GradX_DDD_NT(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, fhs)).cuda().half() +# w = torch.randn((ne, fhs, hs)).cuda().half() +# out = torch.bmm(x, w) +# w = torch.transpose(w, 1, 2) +# +# def benchmark(): +# return torch.bmm(out, w) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::GradX::DDD::NT', +# arguments, +# mean_t, +# std_t, +# x.numel() * hs * 2, +# ) +# +# @parameterized.parameters(*_MATMUL_TESTS) +# def testFFN_Linear1_GradW_DDD_TN(self, sl, hs, fhs, ne): +# assert (sl % ne) == 0 +# x = torch.randn((ne, sl // ne, fhs)).cuda().half() +# w = torch.randn((ne, fhs, hs)).cuda().half() +# out = torch.bmm(x, w) +# x = torch.transpose(x, 1, 2) +# +# def benchmark(): +# return torch.bmm(x, out) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'ffn_hidden_size': fhs, +# 'num_experts': ne, +# } +# log_benchmark( +# '1::GradW::DDD::TN', +# arguments, +# mean_t, +# std_t, +# x.numel() * hs * 2, +# ) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/padded_gather.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/padded_gather.py new file mode 100644 index 0000000000000000000000000000000000000000..c1cf4047c9494394d2a3884ba8830179013db7ff --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/padded_gather.py @@ -0,0 +1,55 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for padded_gather kernel. +class PaddedGatherOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + bins: torch.Tensor, + padded_bins: torch.Tensor, + top_k: int, + ): + ctx.save_for_backward(indices, bin_ids, bins, padded_bins) + ctx.top_k = top_k + return kernels.padded_gather( + x, + indices, + bin_ids, + None, + bins, + padded_bins, + top_k, + ) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + + indices, bin_ids, bins, padded_bins = ctx.saved_tensors + out = kernels.padded_scatter( + grad, + indices, + bin_ids, + None, + bins, + padded_bins, + ctx.top_k, + ) + return out, None, None, None, None, None + + +padded_gather = PaddedGatherOp.apply diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/padded_scatter.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/padded_scatter.py new file mode 100644 index 0000000000000000000000000000000000000000..61e021b81497e472cda5d72bdac557a0ca92d262 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/padded_scatter.py @@ -0,0 +1,98 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for padded_scatter kernel. +class PaddedScatterOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + padded_bins: torch.Tensor, + top_k: int, + ): + maybe_x = [x] if ctx.needs_input_grad[3] else [] + ctx.save_for_backward( + indices, + bin_ids, + weights, + bins, + padded_bins, + *maybe_x, + ) + ctx.top_k = top_k + ctx.x_shape = x.shape + return kernels.padded_scatter( + x, + indices, + bin_ids, + weights, + bins, + padded_bins, + top_k, + ) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + saved_tensors = ctx.saved_tensors + + indices, bin_ids, weights, bins, padded_bins = saved_tensors[:5] + dgrad = None + if ctx.needs_input_grad[0]: + dgrad = kernels.padded_gather( + grad, + indices, + bin_ids, + weights, + bins, + padded_bins, + ctx.top_k, + ) + + wgrad = None + if ctx.needs_input_grad[3]: # need wgrad + x = saved_tensors[-1] + wgrad = kernels.padded_scatter_wgrad( + x, + grad, + indices, + bin_ids, + bins, + padded_bins, + ctx.top_k, + ) + return dgrad, None, None, wgrad, None, None, None, None + + +def padded_scatter( + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + padded_bins: torch.Tensor, + top_k: int, +): + return PaddedScatterOp.apply( + x, + indices, + bin_ids, + weights, + bins, + padded_bins, + top_k, + ) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/padded_scatter_benchmark.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/padded_scatter_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..3cbe4735891446b46f93170c64c23fe63632bf93 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/padded_scatter_benchmark.py @@ -0,0 +1,66 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import torch +# from absl.testing import parameterized + +from .. import benchmark_util, ops + +_PADDED_SCATTER_BENCHMARK = ( + # dMoE-Medium, 8-way EMP. + (1024 * 16, 1024, 8, 4), + # dMoE-Medium, post-all-to-all. + (1024 * 16 * 4, 1024, 8, 1), +) + + +# class PaddedScatterTest(parameterized.TestCase): +# +# @parameterized.parameters(*_PADDED_SCATTER_BENCHMARK) +# def testPaddedScatter(self, sl, hs, ne, top_k): +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# +# # Randomly assign tokens to experts. +# top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(top_expert, ne) +# padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) +# padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# +# # Sample weights for the scatter reduce. +# weights = torch.rand((sl * top_k,)).cuda().half() +# +# # Gather the data to prepare for backwards. +# x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k) +# +# def benchmark(): +# return ops.padded_scatter( +# x, +# indices, +# bin_ids, +# weights, +# bins, +# padded_bins, +# top_k, +# ) +# +# time, std = benchmark_util.benchmark_function(benchmark) +# benchmark_util.log_benchmark( +# 'Padded Scatter', +# { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# 'top_k': top_k, +# }, +# time, +# std, +# ) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/permute_benchmark.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/permute_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..697abddbb3a2082ec4ddd6d94f89f7faabb34b40 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/permute_benchmark.py @@ -0,0 +1,149 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import torch +# from absl.testing import parameterized + +from .. import benchmark_util, ops + +_PERMUTE_TESTS = ( + (16384, 768, 2), + (16384, 768, 4), + (16384, 768, 8), + (16384, 768, 16), + (16384, 768, 32), + (16384, 768, 64), + (16384, 768, 128), + (16384 * 8, 768, 2), + (16384 * 8, 768, 4), + (16384 * 8, 768, 8), + (16384 * 8, 768, 16), + (16384 * 8, 768, 32), + (16384 * 8, 768, 64), + (16384 * 8, 768, 128), +) + + +# class PermuteBenchmark(parameterized.TestCase): +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testBinnedGather(self, sl, hs, ne): +# # NOTE: Capacity factor == 1. +# ec = sl // ne +# +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# top_expert = torch.randint(0, ne, (sl,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(indices, ne) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# +# def benchmark(): +# return ops.binned_gather(x, indices, bins, ec) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('BinnedGather', arguments, mean_t, std_t) +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testBinnedScatter(self, sl, hs, ne): +# # NOTE: Capacity factor == 1. +# ec = sl // ne +# +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# top_expert = torch.randint(0, ne, (sl,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(indices, ne) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# x = ops.binned_gather(x, indices, bins, ec) +# +# def benchmark(): +# return ops.binned_scatter(x, indices, bins) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('BinnedScatter', arguments, mean_t, std_t) +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testPaddedGather(self, sl, hs, ne): +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# +# # Randomly assign tokens to experts. +# top_expert = torch.randint(0, ne, (sl,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(top_expert, ne) +# padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) +# padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# +# def benchmark(): +# return ops.padded_gather(x, indices, bin_ids, bins, padded_bins) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('PaddedGather', arguments, mean_t, std_t) +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testPaddedScatter(self, sl, hs, ne): +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# +# # Randomly assign tokens to experts. +# top_expert = torch.randint(0, ne, (sl,)).cuda().int() +# bin_ids, indices = ops.sort(top_expert) +# tokens_per_expert = ops.histogram(top_expert, ne) +# padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) +# padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) +# bins = ops.inclusive_cumsum(tokens_per_expert, 0) +# x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins) +# +# def benchmark(): +# return ops.padded_scatter(x, indices, bin_ids, bins, padded_bins) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('PaddedScatter', arguments, mean_t, std_t) +# +# @parameterized.parameters(*_PERMUTE_TESTS) +# def testCopy(self, sl, hs, ne): +# # NOTE: Capacity factor == 1. +# # ec = sl // ne +# +# # Create the data and indices. +# x = torch.randn((sl, hs)).cuda().half() +# y = x.clone() +# +# def benchmark(): +# return y.copy_(x) +# +# mean_t, std_t = benchmark_util.benchmark_function(benchmark) +# arguments = { +# 'sequence_length': sl, +# 'hidden_size': hs, +# 'num_experts': ne, +# } +# benchmark_util.log_benchmark('Copy', arguments, mean_t, std_t) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/repeat.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/repeat.py new file mode 100644 index 0000000000000000000000000000000000000000..7e9e09de5f857d51cd758ab30b2f3a846d6f9275 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/repeat.py @@ -0,0 +1,10 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch + + +def repeat(x: torch.Tensor, tiling: torch.Size): + if all((t == 1 for t in tiling)): + return x + return x.repeat(*tiling) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/replicate.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/replicate.py new file mode 100644 index 0000000000000000000000000000000000000000..26daf0eede330603a4b8ea7167faf1411d07ca93 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/replicate.py @@ -0,0 +1,36 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + + +# Autograd wrapper for replicate kernel. +class ReplicateOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, bins: torch.Tensor, num_outputs: int): + ctx.save_for_backward(bins) + out = torch.empty((x.shape[0], num_outputs), dtype=x.dtype, device=x.device) + ops.replicate_forward(x, bins, out) + return out + + @staticmethod + def backward(ctx: Any, grad: torch.Tensor): + bins, = ctx.saved_tensors + out = torch.empty((grad.shape[0], bins.shape[0]), dtype=grad.dtype, device=grad.device) + ops.replicate_backward(grad, bins, out) + return out, None, None + + +replicate = ReplicateOp.apply diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/round_up.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/round_up.py new file mode 100644 index 0000000000000000000000000000000000000000..6cf6bc873c9f448c5fa9126ebcfd66e8688002af --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/round_up.py @@ -0,0 +1,14 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import torch + + +def round_up(x: torch.Tensor, value: int): + assert isinstance(value, int) + assert x.dtype == torch.int32 + + # TODO(tgale): If this becomes and issue + # do this in a custom kernel. We only expect + # to use this on arrays of less than 1k elements. + return torch.div(x + (value - 1), value, rounding_mode='trunc') * value diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/scatter.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/scatter.py new file mode 100644 index 0000000000000000000000000000000000000000..f4605d9b46f387761b070352365f223dbfe69d47 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/scatter.py @@ -0,0 +1,72 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Optional + +import torch +from .stk_autocast import custom_bwd, custom_fwd + +from ..backend import kernels + + +# Autograd wrapper for scatter kernel. +class ScatterOp(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward( + ctx: Any, + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + top_k: int, + ) -> torch.Tensor: + maybe_x = [x] if ctx.needs_input_grad[3] else [] + ctx.save_for_backward(indices, bin_ids, weights, bins, *maybe_x) + ctx.top_k = top_k + ctx.x_shape = x.shape + return kernels.scatter(x, indices, bin_ids, weights, bins, top_k) + + @staticmethod + @custom_bwd + def backward(ctx: Any, grad: torch.Tensor): + grad = grad.contiguous() + saved_tensors = ctx.saved_tensors + + indices, bin_ids, weights, bins = saved_tensors[:4] + dgrad = None + if ctx.needs_input_grad[0]: + dgrad = kernels.gather( + grad, + indices, + bin_ids, + weights, + bins, + ctx.top_k, + ) + + wgrad = None + if ctx.needs_input_grad[3]: # need wgrad + x = saved_tensors[-1] + wgrad = kernels.scatter_wgrad( + x, + grad, + indices, + bin_ids, + bins, + ctx.top_k, + ) + return dgrad, None, None, wgrad, None, None, None + + +def scatter( + x: torch.Tensor, + indices: torch.Tensor, + bin_ids: torch.Tensor, + weights: torch.Tensor, + bins: torch.Tensor, + top_k: int, +) -> Optional[torch.Tensor]: + return ScatterOp.apply(x, indices, bin_ids, weights, bins, top_k) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/sort.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/sort.py new file mode 100644 index 0000000000000000000000000000000000000000..bda3bf64283e39533c2eae3627e76bb2d0262c9f --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/sort.py @@ -0,0 +1,38 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Optional, Tuple + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + +_BITS_FOR_DTYPE = { + torch.int16: 16, + torch.int32: 32, + torch.int64: 64, +} + + +# Autograd wrapper for sort kernel. +# NOTE: Does not support gradients. +class SortOp(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, end_bit: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]: + if end_bit is None: + end_bit = _BITS_FOR_DTYPE[x.dtype] + x_out = torch.empty_like(x) + iota_out = torch.empty_like(x) + ops.sort(x, end_bit, x_out, iota_out) + return (x_out, iota_out) + + +sort = SortOp.apply diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/sort_benchmark.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/sort_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..11043c0824c36372585f1d9f48480c2a6ef32eb6 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/sort_benchmark.py @@ -0,0 +1,85 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import numpy as np +import torch +# from absl.testing import parameterized + +from .. import ops + +_SORT_TESTS = ( + (16384, torch.int32, None), + (16384, torch.int32, 2), + (16384, torch.int32, 128), +) + +_BASELINE_SORT_TESTS = ((16384,),) + + +def numpy_dtype(dtype): + types = { + torch.int16: np.int16, + torch.int32: np.int32, + torch.int64: np.int64, + } + return types[dtype] + + +def benchmark_function(fn, iterations=10): + # Run once to get rid of startup overhead. + fn() + times = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + times = np.array(times) + return times.mean(), times.std(), times.max(), times.min() + + +def log_benchmark(arguments, mean_t, std_t): + print('=' * 60) + print('Benchmark Parameters:') + for (key, value) in arguments.items(): + print(f'{key} = {value}') + print('Results:') + print('mean / std = {:.2f}ms / {:.2f}ms'.format(mean_t, std_t)) + print('=' * 60) + + +# class SortBenchmark(parameterized.TestCase): +# +# @parameterized.parameters(*_SORT_TESTS) +# def testSort(self, n, dtype, max_val): +# if max_val is None: +# max_val = np.iinfo(numpy_dtype(dtype)).max +# end_bit = int(np.ceil(np.log2(max_val))) +# x = torch.randint(0, max_val, (n,)).cuda().to(dtype) +# +# mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.sort(x, end_bit),) +# arguments = { +# 'n': n, +# 'dtype': dtype, +# 'max_val': max_val, +# } +# log_benchmark(arguments, mean_t, std_t) +# +# @parameterized.parameters(*_BASELINE_SORT_TESTS) +# def testTorchSort(self, n): +# x = torch.randint(0, 128, (n,)).cuda().to(torch.int32) +# +# mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.sort(x)) +# arguments = { +# 'n': n, +# } +# log_benchmark(arguments, mean_t, std_t) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/stk_autocast.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/stk_autocast.py new file mode 100644 index 0000000000000000000000000000000000000000..7a3626e5e0eec51339c95a448bca84be14a2ca93 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/stk_autocast.py @@ -0,0 +1,39 @@ +# vendored from +# https://github.com/stanford-futuredata/stk/blob/736313768ef697ce13a0594a41b2512a0fbc9884/stk/backend/autocast.py +import functools +import torch + + +def _is_eligible(x): + return x.is_floating_point() and x.is_cuda and (x.dtype is not torch.float64) + + +def _cast(x, dtype): + if isinstance(x, torch.Tensor) and _is_eligible(x): + return x.to(dtype) + elif isinstance(x, map): + return {_cast(k, dtype): _cast(v, dtype) for k, v in x.items()} + elif isinstance(x, list) or isinstance(x, tuple): + return type(x)(map(lambda y: _cast(y, dtype), x)) + return x + + +def custom_fwd(fwd): + """Wrap a custom autograd function that always uses autocast dtype.""" + + @functools.wraps(fwd) + def decorate_fwd(*args, **kwargs): + if torch.is_autocast_enabled(): + with torch.autocast(device_type="cuda", enabled=False): + dtype = torch.get_autocast_gpu_dtype() + return fwd(*_cast(args, dtype), **_cast(kwargs, dtype)) + return fwd(*args, **kwargs) + return decorate_fwd + + +def custom_bwd(bwd): + @functools.wraps(bwd) + def decorate_bwd(*args, **kwargs): + with torch.autocast(device_type="cuda", enabled=False): + return bwd(*args, **kwargs) + return decorate_bwd \ No newline at end of file diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/sum.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/sum.py new file mode 100644 index 0000000000000000000000000000000000000000..e00c1aa68e1193f5b72f75a2edc37de8d505facc --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/sum.py @@ -0,0 +1,9 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 +import torch + + +def sum(x: torch.Tensor, dim: int = 0): + if x.shape[dim] == 1: + return x.squeeze(dim=dim) + return x.sum(dim=dim) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/ops/topology.py b/build/torch214-cxx11-cu132-x86_64-linux/ops/topology.py new file mode 100644 index 0000000000000000000000000000000000000000..76a50d3164db20534b099dcb4d8487a7aef25d15 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/ops/topology.py @@ -0,0 +1,45 @@ +# Copyright 2024 Databricks +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch + +# Wrap this in a try-block with better error message and +# instructions for building the c++ operations. +try: + from .._ops import ops # type: ignore +except ModuleNotFoundError as e: + raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e + + +# Autograd wrapper for topology kernel. +# NOTE: Does not support gradients. +class TopologyOp(torch.autograd.Function): + + @staticmethod + def forward( + ctx: Any, + padded_bins: torch.Tensor, + block_size: int, + output_block_rows: int, + output_block_columns: int, + ): + out = torch.empty( + output_block_rows * output_block_columns, + dtype=torch.int16, + device=padded_bins.device, + ) + ops.indices( + padded_bins, + block_size, + output_block_rows, + output_block_columns, + out, + ) + return out + + +topology = TopologyOp.apply diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..73c40c267b1c3f4949e9c957a5d2c9f682dfc1a6 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/__init__.py @@ -0,0 +1,7 @@ +# import stk.random +# import stk.ops +# from stk.matrix import Matrix + +from . import random +from . import ops +from .matrix import Matrix diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/backend/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/backend/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/backend/autocast.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/backend/autocast.py new file mode 100644 index 0000000000000000000000000000000000000000..97f6e919a60f3fd579ed0215031008d14111dc96 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/backend/autocast.py @@ -0,0 +1,37 @@ +import functools +import torch + + +def _is_eligible(x): + return x.is_floating_point() and x.is_cuda and (x.dtype is not torch.float64) + + +def _cast(x, dtype): + if isinstance(x, torch.Tensor) and _is_eligible(x): + return x.to(dtype) + elif isinstance(x, map): + return {_cast(k, dtype): _cast(v, dtype) for k, v in x.items()} + elif isinstance(x, list) or isinstance(x, tuple): + return type(x)(map(lambda y: _cast(y, dtype), x)) + return x + + +def custom_fwd(fwd): + """Wrap a custom autograd function that always uses autocast dtype.""" + + @functools.wraps(fwd) + def decorate_fwd(*args, **kwargs): + if torch.is_autocast_enabled(): + with torch.autocast(device_type="cuda", enabled=False): + dtype = torch.get_autocast_gpu_dtype() + return fwd(*_cast(args, dtype), **_cast(kwargs, dtype)) + return fwd(*args, **kwargs) + return decorate_fwd + + +def custom_bwd(bwd): + @functools.wraps(bwd) + def decorate_bwd(*args, **kwargs): + with torch.autocast(device_type="cuda", enabled=False): + return bwd(*args, **kwargs) + return decorate_bwd diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/backend/sputnik.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/backend/sputnik.py new file mode 100644 index 0000000000000000000000000000000000000000..220c947bc1e932e8c77cc30f4069e9930f1aa962 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/backend/sputnik.py @@ -0,0 +1,316 @@ +import torch + +from ..backend import triton_kernels as backend +from ..backend.autocast import custom_bwd, custom_fwd + + +def _standardize_shape(x, transpose): + if transpose: + return torch.Size((x[1], x[0])) + return x + + +def _sparse_transpose(x): + return (torch.Size((x[0][1], x[0][0])), ) + x[1:] + + +def _transpose_helper(x, transpose): + if isinstance(x, torch.Tensor): + return x.t() if transpose else x + if transpose: + x = _sparse_transpose(x) + return x + (transpose,) + + +def _wrap(x): + if isinstance(x, torch.Tensor): + return (x,) + return x + + +def _is_transposed(x): + return (not x.is_contiguous() and + x.stride()[0] == 1 and + x.stride()[1] == x.size()[0]) + + +def _call_helper(op, out, a, b, trans_a, trans_b): + args = (_wrap(_transpose_helper(a, trans_a)) + + _wrap(_transpose_helper(b, trans_b))) + if isinstance(out, tuple): + args = args + out + return op(*args) + + +def _preprocess_inputs(lhs, rhs, dy): + if isinstance(lhs, torch.Tensor) and _is_transposed(lhs): + lhs = lhs.t() + if isinstance(rhs, torch.Tensor) and _is_transposed(rhs): + rhs = rhs.t() + if (isinstance(dy, torch.Tensor) and + not dy.is_contiguous() and + not _is_transposed(dy)): + dy = dy.contiguous() + if isinstance(dy, tuple) and not dy[1].is_contiguous(): + dy = (dy[0], dy[1].contiguous()) + dy[2:] + return lhs, rhs, dy + + +def _postprocess_outputs(x, transpose, grad): + if isinstance(x, torch.Tensor) and transpose: + return grad.t() + return grad + + +def _lhs_gradient(op, lhs, rhs, dy, trans_lhs, trans_rhs): + lhs, rhs, dy = _preprocess_inputs(lhs, rhs, dy) + + a, b = (rhs, dy) if trans_lhs else (dy, rhs) + trans_a = trans_lhs and trans_rhs + trans_b = trans_lhs or not trans_rhs + out = _call_helper(op, lhs, a, b, trans_a, trans_b) + return _postprocess_outputs(lhs, trans_lhs, out) + + +def _rhs_gradient(op, lhs, rhs, dy, trans_lhs, trans_rhs): + lhs, rhs, dy = _preprocess_inputs(lhs, rhs, dy) + + a, b = (dy, lhs) if trans_rhs else (lhs, dy) + trans_a = not trans_lhs or trans_rhs + trans_b = trans_lhs and trans_rhs + out = _call_helper(op, rhs, a, b, trans_a, trans_b) + return _postprocess_outputs(rhs, trans_rhs, out) + + +class DSD(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward(ctx, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_a, + rhs): + ctx.save_for_backward(data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + rhs) + ctx.shape = _standardize_shape(shape, transpose_a) + ctx.transpose_a = transpose_a + + out = torch.empty( + (shape[0], rhs.size()[1]), + dtype=rhs.dtype, + device=rhs.device) + + backend.dsd(shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_a, + rhs, + out) + return out + + @staticmethod + @custom_bwd + def backward(ctx, dy): + saved_tensors = ctx.saved_tensors + lhs = (ctx.shape,) + saved_tensors[:-1] + rhs = saved_tensors[-1] + trans_a = ctx.transpose_a + trans_b = _is_transposed(rhs) + + ddata = None + if ctx.needs_input_grad[1]: + ddata = _lhs_gradient(sdd, + lhs, + rhs, + dy, + trans_a, + trans_b) + drhs = None + if ctx.needs_input_grad[-1]: + op = dds if trans_b else dsd + drhs = _rhs_gradient(op, + lhs, + rhs, + dy, + trans_a, + trans_b) + return None, ddata, None, None, None, None, None, None, None, drhs + + +dsd = DSD.apply + + +class DDS(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward(ctx, + lhs, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_b): + ctx.save_for_backward(lhs, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t) + ctx.shape = _standardize_shape(shape, transpose_b) + ctx.transpose_b = transpose_b + out = torch.empty((lhs.size()[0], shape[1]), + dtype=lhs.dtype, + device=lhs.device) + backend.dds(lhs, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_b, + out) + return out + + @staticmethod + @custom_bwd + def backward(ctx, dy): + saved_tensors = ctx.saved_tensors + lhs = saved_tensors[0] + rhs = (ctx.shape,) + saved_tensors[1:] + trans_a = _is_transposed(lhs) + trans_b = ctx.transpose_b + + dlhs = None + if ctx.needs_input_grad[0]: + op = dsd if trans_a else dds + dlhs = _lhs_gradient(op, + lhs, + rhs, + dy, + trans_a, + trans_b) + ddata = None + if ctx.needs_input_grad[2]: + ddata = _rhs_gradient(sdd, + lhs, + rhs, + dy, + trans_a, + trans_b) + return dlhs, None, ddata, None, None, None, None, None, None, None + + +dds = DDS.apply + + +class SDD(torch.autograd.Function): + + @staticmethod + @custom_fwd + def forward(ctx, + lhs, + rhs, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t): + ctx.save_for_backward( + lhs, + rhs, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t) + ctx.shape = shape + out = torch.empty( + data.shape, + dtype=lhs.dtype, + device=lhs.device) + backend.sdd(lhs, + rhs, + shape, + out, + offsets, + row_indices, + column_indices) + return out + + @staticmethod + @custom_bwd + def backward(ctx, dy): + saved_tensors = ctx.saved_tensors + lhs, rhs = saved_tensors[:2] + dy = (ctx.shape, dy) + saved_tensors[2:] + trans_a = _is_transposed(lhs) + trans_b = _is_transposed(rhs) + + dlhs = None + if ctx.needs_input_grad[0]: + op = dds if trans_a else dsd + dlhs = _lhs_gradient(op, + lhs, + rhs, + dy, + trans_a, + trans_b) + drhs = None + if ctx.needs_input_grad[1]: + op = dsd if trans_b else dds + drhs = _rhs_gradient(op, + lhs, + rhs, + dy, + trans_a, + trans_b) + return dlhs, drhs, None, None, None, None, None, None, None, None + + +sdd = SDD.apply + +class RowIndices(torch.autograd.Function): + + @staticmethod + def forward(ctx, shape, data, offsets, column_indices): + out = torch.empty( + column_indices.shape, + dtype=column_indices.dtype, + device=column_indices.device) + backend.row_indices(shape, data, offsets, column_indices, out) + return out + + +row_indices = RowIndices.apply diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/backend/triton_kernels.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/backend/triton_kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..c535309f3321249f475367164a558f94a4f8eb86 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/backend/triton_kernels.py @@ -0,0 +1,393 @@ +import torch +import triton +import triton.language as tl +from dataclasses import dataclass + +@dataclass +class TritonConfig: + BLOCK_M: int = 128 + BLOCK_N: int = 128 + BLOCK_K: int = 32 + BLOCK_SIZE: int = 128 + NUM_STAGES: int = 4 + NUM_WARPS: int = 4 + +def _validate_matmul_dims(M: int, K: int, N: int): + error_string = "incompatible dimensions: tensor has dim with length: {}, which must be divisible by {}" + assert M % TritonConfig.BLOCK_M == 0, error_string.format(M, TritonConfig.BLOCK_M) + assert K % TritonConfig.BLOCK_K == 0, error_string.format(K, TritonConfig.BLOCK_K) + assert N % TritonConfig.BLOCK_N == 0, error_string.format(N, TritonConfig.BLOCK_N) + +@triton.autotune( + configs=[ + # basic configs for compute-bound matmuls + triton.Config({ + 'BLOCK_M': TritonConfig.BLOCK_M, + 'BLOCK_N': TritonConfig.BLOCK_N, + 'BLOCK_K': TritonConfig.BLOCK_K, + 'BLOCK_SIZE': TritonConfig.BLOCK_SIZE + }, num_stages=TritonConfig.NUM_STAGES, num_warps=TritonConfig.NUM_WARPS), + ], + key=['M', 'N', 'K'], +) +@triton.jit +def _sdd_kernel(A, B, C, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + row_indices, column_indices, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + BLOCK_SIZE: tl.constexpr, GROUP_M: tl.constexpr, ACC_TYPE: tl.constexpr, + ): + # matrix multiplication + pid = tl.program_id(0) + pid_m = tl.load(row_indices + pid) + pid_n = tl.load(column_indices + pid) + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) + rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) + rk = tl.arange(0, BLOCK_K) + # pointers + A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) + B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) + # do matrix multiplication + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) + for k in range(0, tl.cdiv(K, BLOCK_K)): + a = tl.load(A) + b = tl.load(B) + acc += tl.dot(a, b) + A += BLOCK_K * stride_ak + B += BLOCK_K * stride_bk + #Store to sparse matrix + acc = acc.to(C.dtype.element_ty) + BLOCK_ELEMENTS = BLOCK_SIZE * BLOCK_SIZE + cm = tl.arange(0, BLOCK_M) + cn = tl.arange(0, BLOCK_N) + C = C + pid * BLOCK_ELEMENTS + (cm[:, None] * stride_cm + cn[None, :] * stride_cn) + tl.store(C, acc, mask=True) + +@triton.autotune( + configs=[ + # basic configs for compute-bound matmuls + triton.Config({ + 'BLOCK_M': TritonConfig.BLOCK_M, + 'BLOCK_N': TritonConfig.BLOCK_N, + 'BLOCK_K': TritonConfig.BLOCK_K, + 'BLOCK_SIZE': TritonConfig.BLOCK_SIZE + }, num_stages=TritonConfig.NUM_STAGES, num_warps=TritonConfig.NUM_WARPS), + ], + key=['M', 'N', 'K'], +) +@triton.jit +def _dsd_kernel(A, B, C, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + row_indices, column_indices, offsets, + block_offsets_t, trans_A: tl.constexpr, trans_B: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + BLOCK_SIZE: tl.constexpr, GROUP_M: tl.constexpr, ACC_TYPE: tl.constexpr, + ): + + # matrix multiplication + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + num_pid_m = tl.num_programs(0) + num_pid_n = tl.num_programs(1) + pid_n, pid_m = tl.swizzle2d(pid_n, pid_m, num_pid_n, num_pid_m, GROUP_M) + + start_inx = tl.load(offsets + pid_m) + end_inx = tl.load(offsets + pid_m + 1) + + # pointers to sparse matrix + rm = tl.arange(0, BLOCK_M) + rak = tl.arange(0, BLOCK_K) + + A += (rm[:, None] * stride_am + rak[None, :] * stride_ak) + + # pointers to dense matrix + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + rbk = tl.arange(0, BLOCK_K) + B += (rbk[:, None] * stride_bk + rn[None, :] * stride_bn) + + # do matrix multiplication + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) + nsub_blocks = tl.cdiv(BLOCK_SIZE, BLOCK_K) + + BLOCK_ELEMENTS = BLOCK_SIZE * BLOCK_SIZE + ak_sub_incr = BLOCK_K * stride_ak + bk_sub_incr = BLOCK_K * stride_bk + bk_block_incr = BLOCK_SIZE * stride_bk + + for k in range(nsub_blocks * (end_inx - start_inx)): + sub_block_inx = k % nsub_blocks + block_inx = k // nsub_blocks + + if trans_A: + ptr_A = A + tl.load(block_offsets_t + start_inx + block_inx) * BLOCK_ELEMENTS + sub_block_inx * ak_sub_incr + else: + ptr_A = A + (start_inx + block_inx) * BLOCK_ELEMENTS + sub_block_inx * ak_sub_incr + + ptr_B = B + tl.load(column_indices + start_inx + block_inx) * bk_block_incr + sub_block_inx * bk_sub_incr + + a = tl.load(ptr_A) + b = tl.load(ptr_B) + acc += tl.dot(a, b) + + acc = acc.to(C.dtype.element_ty) + + cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + C = C + (cm[:, None] * stride_cm + cn[None, :] * stride_cn) + tl.store(C, acc, mask=True) + +@triton.autotune( + configs=[ + # basic configs for compute-bound matmuls + triton.Config({ + 'BLOCK_M': TritonConfig.BLOCK_M, + 'BLOCK_N': TritonConfig.BLOCK_N, + 'BLOCK_K': TritonConfig.BLOCK_K, + 'BLOCK_SIZE': TritonConfig.BLOCK_SIZE + }, num_stages=TritonConfig.NUM_STAGES, num_warps=TritonConfig.NUM_WARPS), + ], + key=['M', 'N', 'K'], +) +@triton.jit +def _dds_kernel(A, B, C, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + row_indices, column_indices, offsets, + block_offsets_t, trans_A: tl.constexpr, trans_B: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + BLOCK_SIZE: tl.constexpr, GROUP_M: tl.constexpr, ACC_TYPE: tl.constexpr, + ): + + # matrix multiplication + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + num_pid_m = tl.num_programs(0) + num_pid_n = tl.num_programs(1) + pid_n, pid_m = tl.swizzle2d(pid_n, pid_m, num_pid_n, num_pid_m, GROUP_M) + + start_inx = tl.load(offsets + pid_n) + end_inx = tl.load(offsets + pid_n + 1) + + # pointers to dense matrix + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rak = tl.arange(0, BLOCK_K) + + A += (rm[:, None] * stride_am + rak[None, :] * stride_ak) + + # pointers to sparse matrix + rn = tl.arange(0, BLOCK_N) + rbk = tl.arange(0, BLOCK_K) + B += (rbk[:, None] * stride_bk + rn[None, :] * stride_bn) + + # do matrix multiplication + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) + nsub_blocks = tl.cdiv(BLOCK_SIZE, BLOCK_K) + + BLOCK_ELEMENTS = BLOCK_SIZE * BLOCK_SIZE + + ak_sub_incr = BLOCK_K * stride_ak + ak_block_incr = BLOCK_SIZE * stride_ak + bk_sub_incr = BLOCK_K * stride_bk + + for k in range(nsub_blocks * (end_inx - start_inx)): + sub_block_inx = k % nsub_blocks + block_inx = k // nsub_blocks + + if trans_B: + ptr_B = B + (start_inx + block_inx) * BLOCK_ELEMENTS + sub_block_inx * bk_sub_incr + else: + ptr_B = B + tl.load(block_offsets_t + start_inx + block_inx) * BLOCK_ELEMENTS + sub_block_inx * bk_sub_incr + + ptr_A = A + tl.load(column_indices + start_inx + block_inx) * ak_block_incr + sub_block_inx * ak_sub_incr + a = tl.load(ptr_A) + b = tl.load(ptr_B) + acc += tl.dot(a, b) + + acc = acc.to(C.dtype.element_ty) + cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + C = C + (cm[:, None] * stride_cm + cn[None, :] * stride_cn) + tl.store(C, acc, mask=True) + +def dsd(shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_a, + rhs, + out + ): + + device = rhs.device + trans_A = transpose_a + trans_B = False + + if rhs.stride(0) > 1 and rhs.stride(1) > 1: + trans_B = True + + # checks constraints + assert shape[1] == rhs.shape[0], "incompatible dimensions" + M, K = shape + _, N = rhs.shape + + _validate_matmul_dims(M, K, N) + + # accumulator types + ACC_TYPE = tl.float32 if rhs.dtype in [torch.float16, torch.bfloat16, torch.float32] else tl.int32 + + stride_am, stride_ak = data.stride(1), data.stride(2) + stride_bk, stride_bn = rhs.stride(0), rhs.stride(1) + a_column_indices = column_indices + a_offsets = offsets + + # launch kernel + grid = lambda META: (triton.cdiv(M, META['BLOCK_M']), triton.cdiv(N, META['BLOCK_N'])) + + if trans_A: + stride_am, stride_ak = data.stride(2), data.stride(1) + a_column_indices, a_offsets = column_indices_t, offsets_t + + if trans_B: + stride_bk, stride_bn = rhs.stride(1), rhs.stride(0) + + _dsd_kernel[grid]( + data.data, rhs, out, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + out.stride(0), out.stride(1), + row_indices, a_column_indices, a_offsets, + block_offsets_t, trans_A, trans_B, + GROUP_M=128, ACC_TYPE=ACC_TYPE + ) + # return out + +def dds(lhs, + shape, + data, + offsets, + row_indices, + column_indices, + offsets_t, + column_indices_t, + block_offsets_t, + transpose_b, + out + ): + + device = lhs.device + trans_B = transpose_b + trans_A = False + + if lhs.stride(0) > 1 and lhs.stride(1) > 1: + trans_A = True + + # checks constraints + assert lhs.shape[1] == shape[0], "incompatible dimensions" + M, K = lhs.shape + _, N = shape + + _validate_matmul_dims(M, K, N) + + # accumulator types + ACC_TYPE = tl.float32 if lhs.dtype in [torch.float16, torch.bfloat16, torch.float32] else tl.int32 + + stride_am, stride_ak = lhs.stride(0), lhs.stride(1) + stride_bk, stride_bn = data.stride(1), data.stride(2) + b_column_indices = column_indices_t + b_offsets = offsets_t + + # launch kernel + grid = lambda META: (triton.cdiv(M, META['BLOCK_M']), triton.cdiv(N, META['BLOCK_N'])) + + if trans_A: + stride_am, stride_ak = lhs.stride(1), lhs.stride(0) + if trans_B: + stride_bk, stride_bn = data.stride(2), data.stride(1) + b_column_indices, b_offsets = column_indices, offsets + + _dds_kernel[grid]( + lhs, data, out, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + out.stride(0), out.stride(1), + row_indices, b_column_indices, b_offsets, + block_offsets_t, trans_A, trans_B, + GROUP_M=128, ACC_TYPE=ACC_TYPE + ) + +def sdd(lhs, + rhs, + shape, + out, + offsets, + row_indices, + column_indices + ): + + device = out.device + trans_A = False + trans_B = False + + if lhs.stride(0) > 1 and lhs.stride(1) > 1: + trans_A = True + if rhs.stride(0) > 1 and rhs.stride(1) > 1: + trans_B = True + + # checks constraints + assert lhs.shape[1] == rhs.shape[0], "incompatible dimensions" + M, K = lhs.shape + _, N = rhs.shape + + _validate_matmul_dims(M, K, N) + + # accumulator types + ACC_TYPE = tl.float32 if out.dtype in [torch.float16, torch.bfloat16, torch.float32] else tl.int32 + + # launch kernel + nnz_blocks = len(row_indices) + grid = lambda META: (nnz_blocks,) + + stride_am, stride_ak = lhs.stride(0), lhs.stride(1) + stride_bk, stride_bn = rhs.stride(0), rhs.stride(1) + + if trans_A: + stride_am, stride_ak = lhs.stride(1), lhs.stride(0) + if trans_B: + stride_bk, stride_bn = rhs.stride(1), rhs.stride(0) + + _sdd_kernel[grid]( + lhs, rhs, out, M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + out.stride(1), out.stride(2), + row_indices, column_indices, + GROUP_M=128, ACC_TYPE=ACC_TYPE + ) + +@triton.jit +def _row_indices_kernel(offsets, out): + pid = tl.program_id(0) + row_offset = tl.load(offsets + pid) + nnz_blocks = tl.load(offsets + pid + 1) - row_offset + for nnz_block in range(nnz_blocks): + tl.store(out + row_offset + nnz_block, pid) + +def row_indices( + shape, data, offsets, column_indices, out +): + block_rows = len(offsets) - 1 + _row_indices_kernel[(block_rows, )](offsets, out) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/matrix.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/matrix.py new file mode 100644 index 0000000000000000000000000000000000000000..80f42263d6aada287adbfa52a61fe950162a9e28 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/matrix.py @@ -0,0 +1,329 @@ +import numpy as np +import torch + +# 1. Add heavyweight (data) validation helper. +# 2. Add construction helpers +# 3. Make indentation consistent +# 4. Replace asserts with descriptive errors. + +## +### Validation helpers. +## + + +def _validate_matrix(shape, data, row_indices, column_indices, offsets): + # Data should be [nnz, block_size, block_size] + if data.dim() == 1: + data = torch.reshape(data, [data.numel(), 1, 1]) + + # Blocks should be square. + if data.shape[-2] != data.shape[-1]: + raise ValueError( + "Expected square blocking in data. " + f"Got block shape {[data.shape[-2], data.shape[-1]]}") + + # Flatten batch dimensions on data - original shape preserved + # in shape argument. + block_size = data.shape[-1] + data = data.view([-1, block_size, block_size]) + + if data.dim() != 3: + raise ValueError( + "Expected 3D shape for data (nnz, block, block). " + f"Got shape {data.dim()}D shape.") + + block_size = data.shape[1] + if shape[-2] % block_size != 0 or shape[-1] % block_size != 0: + raise ValueError( + "Matrix shape must be dividible by blocking. " + f"Got shape {shape} with " + f"{[block_size, block_size]} blocking.") + + if np.prod(shape) < data.numel(): + raise ValueError( + "Invalid matrix. Number of nonzeros exceeds matrix capacity " + f"({data.numel()} v. {np.prod(shape)})") + + if row_indices.dim() != 1: + raise ValueError( + f"Expected 1D row_indices. Got {row_indices.dim()}D row_indices.") + + if column_indices.dim() != 1: + raise ValueError( + f"Expected 1D column_indices. Got {column_indices.dim()}D column_indices.") + + if offsets.dim() != 1: + raise ValueError( + f"Expected 1D offsets. Got {offsets.dim()}D offsets.") + + if row_indices.numel() != data.shape[0]: + raise ValueError( + "Expected 1 index per nonzero block. " + f"Got {row_indices.numel()} row_indices for {data.shape[0]} blocks") + + if column_indices.numel() != data.shape[0]: + raise ValueError( + "Expected 1 index per nonzero block. " + f"Got {column_indices.numel()} column_indices for {data.shape[0]} blocks") + + block_rows = np.prod(shape[:-1]) / block_size + if offsets.numel() != block_rows + 1: + raise ValueError( + "Expected one offset per block row plus one. " + f"Got {offsets.numel()} offsets with {block_rows} block rows.") + + is_cuda = (data.is_cuda and + row_indices.is_cuda and + column_indices.is_cuda and + offsets.is_cuda) + is_cpu = (not data.is_cuda and + not row_indices.is_cuda and + not column_indices.is_cuda and + not offsets.is_cuda) + if not (is_cuda or is_cpu): + raise ValueError( + "Expected data & meta-data on common device. " + f"Got data on {data.device}, row_indices on {row_indices.device} " + f"column_indices on {column_indices.device} and " + f"offsets on {offsets.device}.") + + if data.dtype != torch.float16: + raise ValueError( + f"Expected float16 data. Got {data.dtype} data.") + if row_indices.dtype != torch.int16: + raise ValueError( + f"Expected int16 row_indices. Got {row_indices.dtype} row_indices.") + if column_indices.dtype != torch.int16: + raise ValueError( + f"Expected int16 column_indices. Got {column_indices.dtype} column_indices.") + if offsets.dtype != torch.int32: + raise ValueError( + f"Expected int32 offsets. Got {offsets.dtype} offsets.") + return data + + +def _transpose(size, data, row_indices, column_indices, offsets): + block_columns = size[1] // data.shape[1] + + # Sort row indices by column indices to get the transposed matrix's + # column indices. + gather_indices = column_indices.argsort() + column_indices_t = row_indices.gather(0, gather_indices) + block_offsets_t = gather_indices.int() + + # NOTE: Histogram is not implemented for any integer type on CPU. Do + # the histogram in 32-bit float, which can exactly represent 16-bit + # integers. + column_indices_float = column_indices.float() + + zero = torch.zeros((1,), dtype=torch.int32, device=data.device) + nnz_per_column = column_indices_float.histc(block_columns, 0, block_columns) + nnz_per_column = nnz_per_column.int() + offsets_t = torch.cat([zero, nnz_per_column.cumsum(0, dtype=torch.int32)]) + return column_indices_t, offsets_t, block_offsets_t + + +class Matrix(torch.nn.Module): + """A matrix stored in sparse format. + + Underlying format is block compressed sparse row (BCSR). + + TODO(tgale): Make this mirror torch.Tensor API as much as possible. + """ + + def __init__(self, + size, + data, + row_indices, + column_indices, + offsets, + column_indices_t=None, + offsets_t=None, + block_offsets_t=None): + super().__init__() + self._size = size + self._data = data + self._row_indices = row_indices + self._column_indices = column_indices + self._offsets = offsets + + # Produce the transpose meta-data if it is not passed in. + if ((column_indices_t is None) or (offsets_t is None) or + (block_offsets_t is None)): + column_indices_t, offsets_t, block_offsets_t = _transpose( + size, data, row_indices, column_indices, offsets) + self._column_indices_t = column_indices_t + self._offsets_t = offsets_t + self._block_offsets_t = block_offsets_t + + self._transposed = False + + # Validate that our metadata will not overflow. + max_dim = np.iinfo(np.int16).max * self.blocking + if column_indices.dtype == torch.int16: + if size[0] > max_dim or size[1] > max_dim: + raise ValueError( + "Sparse matrix with shape {size} exceeds representable " + "size with 16-bit indices.") + + def validate(self): + _validate_matrix(self._size, + self._data, + self._row_indices, + self._column_indices, + self._offsets) + + # TODO(tgale): Add heavyweight data validation. + + def to(self, device): + # TODO(tgale): Handle type conversions here. We + # need to set the appropriate meta-data type for + # the given floating-point type. + self._data = self._data.to(device) + self._row_indices = self._row_indices.to(device) + self._column_indices = self._column_indices.to(device) + self._offsets = self._offsets.to(device) + self._column_indices_t = self._column_indices_t.to(device) + self._offsets_t = self._offsets_t.to(device) + self._block_offsets_t = self._block_offsets_t.to(device) + return self + + def cuda(self): + return self.to(torch.cuda.current_device()) + + def clone(self): + return Matrix( + self.size(), + self.data.clone(), + self.row_indices.clone(), + self.column_indices.clone(), + self.offsets.clone(), + self.column_indices_t.clone(), + self.offsets_t.clone(), + self.block_offsets_t.clone()) + + def t(self): + if self.dim() != 2: + raise ValueError( + "t() expects a tensor with <= 2 dimensions, " + f"but self is {self.dim()}D.") + out = Matrix(self.size(), + self.data, + self.row_indices, + self.column_indices, + self.offsets, + self.column_indices_t, + self.offsets_t, + self.block_offsets_t) + out._transposed = not self._transposed + out._size = torch.Size((self._size[1], self._size[0])) + return out + + def contiguous(self): + raise ValueError("Not yet implemented.") + + def is_contiguous(self): + return not self._transposed + + @property + def is_cuda(self): + return self._data.is_cuda + + @property + def device(self): + return self._data.device + + def size(self): + return self._size + + @property + def shape(self): + return self.size() + + def dim(self): + return len(self._size) + + @property + def data(self): + return self._data + + @property + def row_indices(self): + return self._row_indices + + @property + def column_indices(self): + return self._column_indices + + @property + def offsets(self): + return self._offsets + + @property + def offsets_t(self): + return self._offsets_t + + @property + def column_indices_t(self): + return self._column_indices_t + + @property + def block_offsets_t(self): + return self._block_offsets_t + + @property + def dtype(self): + return self.data.dtype + + @property + def nnz(self): + return self.data.numel() + + @property + def blocking(self): + return self.data.shape[1] + + @property + def requires_grad(self): + return self.data.requires_grad + + def requires_grad_(self, x): + self.data.requires_grad_(x) + return self + + def view(self, *shape): + assert self.is_contiguous() + if shape[-1] != self.size()[-1]: + raise ValueError( + "Can't change view on compressed dimension. " + f"{self.size()[-1]} v. {shape[-1]}.") + if np.prod(shape) != np.prod(self.size()): + raise ValueError( + "Mismatch in numel of Matrix and new shape. " + f"{np.prod(self.size())} v. {np.prod(shape)}") + return Matrix(shape, + self.data, + self.row_indices, + self.column_indices, + self.offsets, + self.column_indices_t, + self.offsets_t, + self.block_offsets_t) + + @property + def grad(self): + # TODO(tgale): Make sure this mirrors torch.Tensor + # behavior in the case where we ask for the gradient + # of a non-contiguous tensor. + size = self.size() + if not self.is_contiguous(): + size = torch.Size((size[1], size[0])) + out = Matrix(size, + self.data.grad, + self.row_indices, + self.column_indices, + self.offsets, + self.column_indices_t, + self.offsets_t, + self.block_offsets_t) + return out if self.is_contiguous() else out.t() diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fc873b236f4cd4036964c016a4036e3ce5ebf1ac --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/__init__.py @@ -0,0 +1,3 @@ +from .linear_ops import dds, dsd, sdd +from .matrix_ops import ones_like, row_indices, sum, to_dense, to_sparse +from .eltwise_ops import mul diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/eltwise_ops.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/eltwise_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ba7d7332320250fd01fa60e60528f19de3e8ed03 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/eltwise_ops.py @@ -0,0 +1,28 @@ +from ..matrix import Matrix + +def mul(a, b): + """Performs element-wise multiplication of matrices a and b. + + It is the user's responsibility to make sure that a and b + follow the same matrix topology. This function assumes it is safe + to use the topoplogy of a. + + Args: + a: stk.Matrix. + b: stk.Matrix with a's matrix topology. + + Returns: + stk.Matrix where the entries correspond to torch.mul(a, b). + """ + assert isinstance(a, Matrix) + assert isinstance(b, Matrix) + assert a.size() == b.size() + + return Matrix(a.size(), + a.data * b.data, + a.row_indices, + a.column_indices, + a.offsets, + a.column_indices_t, + a.offsets_t, + a.block_offsets_t) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/eltwise_ops_test.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/eltwise_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2939372a5c68ac92b47b11015db4f75f4fd60ffa --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/eltwise_ops_test.py @@ -0,0 +1,86 @@ +import unittest +import itertools +import torch +# from absl.testing import parameterized + +import stk +from stk.ops.linear_ops_test import allclose, _dense_and_sparse + +_MATRIX_SIZES = ( + (128, 128, 0.0), + (256, 256, 0.5), + (2048, 1024, 0.8), + (512, 128, 0.0), + (128, 512, 0.0), + (1024, 512, 0.0), + (1024, 512, 0.5), + (1024, 512, 0.75), + (512, 1024, 0.0), + (512, 1024, 0.5), + (512, 1024, 0.75), + (1024, 1024, 0.0), + (1024, 1024, 0.5), + (1024, 1024, 0.75), +) + +_DTYPE = ( + torch.float16, torch.bfloat16 +) + +def _generate_testcases(): + testcases = itertools.product(_MATRIX_SIZES, _DTYPE) + testcases = [(*size, 128, dtype) for + (size, dtype) in testcases] + return testcases + +_ELTWISE_OP_TESTS = _generate_testcases() + +def _dense_and_sparse_like(x, std=0.1): + dense_data = torch.randn_like(x.data, device=x.device) * std + sparse = stk.Matrix(x.size(), + dense_data, + x.row_indices, + x.column_indices, + x.offsets) + dense = stk.ops.to_dense(sparse) + + return (dense.requires_grad_(True), + sparse.requires_grad_(True)) + +# @parameterized.parameters(_ELTWISE_OP_TESTS) +# class EltwiseOpsTest(parameterized.TestCase): +# +# def testEltwiseMul(self, m, n, sparsity, blocking, dtype): +# +# a_dense, a = _dense_and_sparse(m, n, sparsity, blocking, dtype) +# b_dense, b = _dense_and_sparse_like(a) +# +# out = stk.ops.mul(a, b) +# expected_out = torch.mul(a_dense, b_dense) +# +# # Compute the gradients w.r.t. the inputs. +# expected_out.sum().backward() +# stk.ops.sum(out).backward() +# +# # Validate the results. +# out = stk.ops.to_dense(out) +# self.assertEqual(out.dim(), 2) +# self.assertEqual(expected_out.size(), out.size()) +# self.assertTrue(allclose(out, expected_out)) +# +# # LHS gradient. +# grad = stk.ops.to_dense(a.grad) +# expected_grad = a_dense.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size(), grad.size()) +# self.assertTrue(allclose(grad, expected_grad)) +# +# # RHS gradient. +# grad = stk.ops.to_dense(b.grad) +# expected_grad = b_dense.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size(), grad.size()) +# self.assertTrue(allclose(grad, expected_grad)) + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/linear_ops.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/linear_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..9d277c8c07f9e30addc31900a12175c8a1f4d7ad --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/linear_ops.py @@ -0,0 +1,59 @@ +import torch + +from ..backend import sputnik +from ..matrix import Matrix + + +def dsd(a, b): + assert isinstance(a, Matrix) + assert isinstance(b, torch.Tensor) + return sputnik.dsd( + a.size(), + a.data, a.offsets, + a.row_indices, + a.column_indices, + a.offsets_t, + a.column_indices_t, + a.block_offsets_t, + not a.is_contiguous(), + b) + + +def dds(a, b): + assert isinstance(a, torch.Tensor) + assert isinstance(b, Matrix) + return sputnik.dds( + a, + b.size(), + b.data, b.offsets, + b.row_indices, + b.column_indices, + b.offsets_t, + b.column_indices_t, + b.block_offsets_t, + not b.is_contiguous()) + + +def sdd(a, b, topo): + assert isinstance(a, torch.Tensor) + assert isinstance(b, torch.Tensor) + assert isinstance(topo, Matrix) + assert topo.is_contiguous() + out = sputnik.sdd( + a, b, + topo.size(), + topo.data, + topo.offsets, + topo.row_indices, + topo.column_indices, + topo.offsets_t, + topo.column_indices_t, + topo.block_offsets_t) + return Matrix(topo.size(), + out, + topo.row_indices, + topo.column_indices, + topo.offsets, + topo.column_indices_t, + topo.offsets_t, + topo.block_offsets_t) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/linear_ops_test.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/linear_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..c1c24d350df9c1b2346c7da885502cd696c88867 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/linear_ops_test.py @@ -0,0 +1,216 @@ +import unittest +import itertools +import numpy as np +import torch +# from absl.testing import parameterized + +import stk + + +def allclose(x, y, pct=0.25): + mask = torch.isclose(x, y, rtol=5e-2) + pct_diff = (mask.numel() - mask.sum()) / mask.numel() * 100 + if pct_diff > pct: + print("{:.2f}% of values not close.".format(pct_diff)) + return False + return True + + +# An assortment of problems designed to make sure +# the bindings are operating correctly. +_MATRIX_SIZES = ( + (128, 128, 128, 0.0), + (256, 256, 256, 0.5), + (2048, 1024, 512, 0.8), + (512, 128, 128, 0.0), + (128, 128, 512, 0.0), + (1024, 512, 512, 0.0), + (1024, 512, 512, 0.5), + (1024, 512, 512, 0.75), + (512, 512, 1024, 0.0), + (512, 512, 1024, 0.5), + (512, 512, 1024, 0.75), + (1024, 1024, 1024, 0.0), + (1024, 1024, 1024, 0.5), + (1024, 1024, 1024, 0.75), +) + +_TRANSPOSE = ( + (False, False), + (False, True), + (True, False), + (True, True), +) + +_DTYPE = ( + torch.float16, torch.bfloat16 +) + +def _generate_testcases(): + testcases = itertools.product(_MATRIX_SIZES, _TRANSPOSE, _DTYPE) + testcases = [(*size, *trans, 128, dtype) for + (size, trans, dtype) in testcases] + return testcases + +_LINEAR_OP_TESTS = _generate_testcases() + +def _dense_and_sparse(rows, cols, sparsity, blocking, dtype, std=0.1): + mask = stk.random.dense_mask(rows, cols, sparsity, blocking) + dense = (torch.randn(rows, cols) * std * mask).type(dtype) + sparse = stk.ops.to_sparse(dense, blocking) + cuda_device = torch.device("cuda") + return (dense.to(cuda_device).requires_grad_(True), + sparse.to(cuda_device).requires_grad_(True)) + + +def _dense(rows, cols, dtype, std=0.1): + cuda_device = torch.device("cuda") + out = (torch.randn(rows, cols) * std).type(dtype) + return out.to(cuda_device).requires_grad_(True) + + +def _dense_2x(rows, cols, dtype): + a = _dense(rows, cols, dtype) + return a, a.detach().requires_grad_(True) + + +def _with_transpose(op, a, b, trans_a, trans_b): + a = a.t() if trans_a else a + b = b.t() if trans_b else b + return op(a, b) + + +def _mmm(a, b, topo): + mask = stk.ops.to_dense(stk.ops.ones_like(topo)) + return torch.mm(a, b) * mask + + +def _sparse_out_with_transpose(op, a, b, topo, trans_a, trans_b): + a = a.t() if trans_a else a + b = b.t() if trans_b else b + return op(a, b, topo) + + +def _mask(x, mask): + mask = stk.ops.to_dense(stk.ops.ones_like(mask)) + return x * mask + + +# @parameterized.parameters(*_LINEAR_OP_TESTS) +# class LinearOpsTest(parameterized.TestCase): +# +# def testLinearOps_Dsd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype): +# # Construct the operands. +# a_shape = (k, m) if trans_a else (m, k) +# a_dense, a = _dense_and_sparse(*a_shape, sparsity, blocking, dtype) +# b_shape = (n, k) if trans_b else (k, n) +# b, bcp = _dense_2x(*b_shape, dtype) +# +# # Execute the matmul. +# out = _with_transpose(stk.ops.dsd, a, b, trans_a, trans_b) +# expected_out = _with_transpose(torch.mm, a_dense, bcp, trans_a, trans_b) +# +# # Compute the gradients w.r.t. the inputs. +# expected_out.sum().backward() +# out.sum().backward() +# +# # Validate the results. +# self.assertEqual(out.dim(), 2) +# self.assertEqual(expected_out.size()[0], out.size()[0]) +# self.assertEqual(expected_out.size()[1], out.size()[1]) +# self.assertTrue(allclose(out, expected_out)) +# +# # LHS gradient. +# grad = stk.ops.to_dense(a.grad) +# expected_grad = _mask(a_dense.grad, a.grad) +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# # RHS gradient. +# grad = b.grad +# expected_grad = bcp.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# def testLinearOps_Dds(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype): +# # Construct the operands. +# a_shape = (k, m) if trans_a else (m, k) +# a, acp = _dense_2x(*a_shape, dtype) +# b_shape = (n, k) if trans_b else (k, n) +# b_dense, b = _dense_and_sparse(*b_shape, sparsity, blocking, dtype) +# +# # Execute the matmul. +# out = _with_transpose(stk.ops.dds, a, b, trans_a, trans_b) +# expected_out = _with_transpose(torch.mm, acp, b_dense, trans_a, trans_b) +# +# # Compute the gradients w.r.t. the inputs. +# expected_out.sum().backward() +# out.sum().backward() +# +# # Validate the results. +# self.assertEqual(out.dim(), 2) +# self.assertEqual(expected_out.size()[0], out.size()[0]) +# self.assertEqual(expected_out.size()[1], out.size()[1]) +# self.assertTrue(allclose(out, expected_out)) +# +# # LHS gradient. +# grad = a.grad +# expected_grad = acp.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# # RHS gradient. +# grad = stk.ops.to_dense(b.grad) +# expected_grad = _mask(b_dense.grad, b.grad) +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# def testLinearOps_Sdd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype): +# # Construct the operands. +# a_shape = (k, m) if trans_a else (m, k) +# a, acp = _dense_2x(*a_shape, dtype) +# b_shape = (n, k) if trans_b else (k, n) +# b, bcp = _dense_2x(*b_shape, dtype) +# _, topo = _dense_and_sparse(m, n, sparsity, blocking, dtype) +# +# # Execute the matmul. +# out = _sparse_out_with_transpose(stk.ops.sdd, a, b, topo, trans_a, trans_b) +# expected_out = _sparse_out_with_transpose(_mmm, acp, bcp, topo, trans_a, trans_b) +# +# # Compute the gradients w.r.t. the inputs. +# expected_out.sum().backward() +# stk.ops.sum(out).backward() +# +# # Validate the results. +# out = stk.ops.to_dense(out) +# self.assertEqual(out.dim(), 2) +# self.assertEqual(expected_out.size()[0], out.size()[0]) +# self.assertEqual(expected_out.size()[1], out.size()[1]) +# self.assertTrue(allclose(out, expected_out)) +# +# # LHS gradient. +# grad = a.grad +# expected_grad = acp.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) +# +# # RHS gradient. +# grad = b.grad +# expected_grad = bcp.grad +# self.assertEqual(grad.dim(), 2) +# self.assertEqual(expected_grad.size()[0], grad.size()[0]) +# self.assertEqual(expected_grad.size()[1], grad.size()[1]) +# self.assertTrue(allclose(grad, expected_grad)) + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/matrix_ops.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/matrix_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..447c72dc73439d84f58c917676cc04e64f13e97d --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/matrix_ops.py @@ -0,0 +1,98 @@ +from ..backend import sputnik +from ..matrix import Matrix +import torch +import numpy as np + + +@torch.no_grad() +def row_indices(shape, data, offsets, column_indices): + return sputnik.row_indices(shape, data, offsets, column_indices) + + +# TODO(tgale): Replace this helper with a custom kernel. This operation +# is much simpler to do than how it's currently implemented. +@torch.no_grad() +def _expand_for_blocking(idxs, blocking): + # Duplicate for block column dimension. + idxs = torch.reshape(idxs, [idxs.size()[0], 1, 2]).repeat(1, blocking, 1) + + # Update the column indices. + idxs[:, :, 1] *= blocking + idxs[:, :, 1] += torch.reshape(torch.arange(blocking, device=idxs.device), [1, blocking]) + + # Duplicate for block row dimension. + idxs = torch.reshape(idxs, [idxs.size()[0], 1, blocking, 2]) + idxs = idxs.repeat(1, blocking, 1, 1) + + # Update the row indices. + idxs[:, :, :, 0] *= blocking + idxs[:, :, :, 0] += torch.reshape(torch.arange(blocking, device=idxs.device), [1, blocking, 1]) + idxs = torch.reshape(idxs, [-1, 2]) + return idxs + + +# TODO(tgale): Add input type checking. +@torch.no_grad() +def to_dense(x): + assert isinstance(x, Matrix) + + shape = (np.prod(x.shape[:-1]), x.shape[-1]) + row_idxs = x.row_indices.type(torch.int32) + col_idxs = x.column_indices.type(torch.int32) + indices = _expand_for_blocking(torch.stack([row_idxs, col_idxs], dim=1), x.blocking) + indices = (indices[:, 0] * shape[1] + indices[:, 1]).type(torch.int64) + + out = torch.zeros(shape[0] * shape[1], dtype=x.dtype, device=x.device) + out.scatter_(0, indices, x.data.flatten()) + return out.reshape(x.size()) + + +@torch.no_grad() +def _mask(x, blocking=1): + assert x.dim() == 2 + assert x.size()[0] % blocking == 0 + assert x.size()[1] % blocking == 0 + block_rows = x.size()[0] // blocking + block_cols = x.size()[1] // blocking + x = torch.reshape(x, [block_rows, blocking, block_cols, blocking]) + x = torch.sum(torch.abs(x), dim=(1, 3)) + return x != 0 + + +# TODO(tgale): Add input type checking. +@torch.no_grad() +def to_sparse(x, blocking=1): + m = _mask(x, blocking) + + # TODO(tgale): Set to appropriate type for input matrix. + row_nnzs = torch.sum(m, dim=1).type(torch.int32) + zeros = torch.zeros((1,), dtype=row_nnzs.dtype, device=row_nnzs.device) + offsets = torch.cat([zeros, torch.cumsum(row_nnzs, dim=0)]) + offsets = offsets.type(torch.int32) + + indices = torch.nonzero(m).type(torch.int16) + row_indices = indices[:, 0] + column_indices = indices[:, 1] + + # Nonzero indices in the dense matrix. + nonzero_indices = torch.nonzero(m) + nonzero_indices = _expand_for_blocking(nonzero_indices, blocking) + nonzero_indices = nonzero_indices[:, 0] * x.size()[1] + nonzero_indices[:, 1] + + # Gather the data and construct the sparse matrix. + data = torch.gather(x.flatten(), dim=0, index=nonzero_indices) + data = torch.reshape(data, [-1, blocking, blocking]) + return Matrix(x.size(), data, row_indices, column_indices, offsets) + + +@torch.no_grad() +def ones_like(x): + return Matrix(x.size(), + torch.ones_like(x.data), + x.row_indices, + x.column_indices, x.offsets) + + +def sum(x): + assert isinstance(x, Matrix) + return x.data.sum() diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/matrix_ops_test.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/matrix_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..d172d921f6f08b0e4fb709207a458b0e1e071bd0 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/ops/matrix_ops_test.py @@ -0,0 +1,62 @@ +import unittest + +# from absl.testing import parameterized +import stk +import torch + + +# @parameterized.parameters( +# (8, 16, 0.0, 1), +# (8, 16, 0.5, 1), +# (8, 16, .95, 1), +# (16, 8, 0.0, 1), +# (16, 8, 0.5, 1), +# (16, 8, .95, 1), +# (8, 16, 0.0, 8), +# (8, 16, 0.5, 8), +# (8, 16, 1.0, 8), +# (16, 8, 0.0, 8), +# (16, 8, 0.5, 8), +# (16, 8, 1.0, 8), +# (128, 256, 0.5, 16), +# (256, 128, 0.75, 32), +# (512, 512, .875, 128)) +# class MatrixOpsTest(parameterized.TestCase): +# +# def testMatrixOps_FormatConversion(self, rows, cols, sparsity, blocking): +# mask = stk.random.dense_mask(rows, cols, sparsity, blocking) +# x = (torch.randn(rows, cols) * mask).type(torch.float16) +# +# # Convert the matrix to sparse format. +# sparse_x = stk.ops.to_sparse(x, blocking) +# +# # Validate the matrix. +# sparse_x.validate() +# +# # Validate the shape. +# self.assertEqual(sparse_x.dim(), 2) +# self.assertEqual(sparse_x.size()[0], rows) +# self.assertEqual(sparse_x.size()[1], cols) +# +# # Validate the sparsity. +# numblocks = rows // blocking * cols // blocking +# nnz = round(numblocks * (1 - sparsity)) * blocking ** 2 +# self.assertEqual(sparse_x.nnz, nnz) +# +# # Convert back to dense format. +# dense_x = stk.ops.to_dense(sparse_x) +# +# # Validate the shape. +# self.assertEqual(dense_x.dim(), 2) +# self.assertEqual(dense_x.size()[0], rows) +# self.assertEqual(dense_x.size()[1], cols) +# +# # Validate the sparsity +# self.assertEqual(torch.count_nonzero(dense_x).item(), nnz) +# +# # Validate the output. +# self.assertTrue(torch.all(torch.eq(x, dense_x))) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/random/__init__.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/random/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2576d1ca27283f77569a9a620c7c99fa68aaf30e --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/random/__init__.py @@ -0,0 +1,2 @@ +# from stk.random.random_ops import dense_mask, mask, randn +from .random_ops import dense_mask, mask, randn diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/random/random_ops.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/random/random_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..d1b36771e0eb8e7abf46bcb3b136b5fb1d29df93 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/random/random_ops.py @@ -0,0 +1,36 @@ +import numpy as np +import torch +from ..ops import matrix_ops + + +@torch.no_grad() +def dense_mask(rows, cols, sparsity, blocking=1): + assert sparsity >= 0.0 and sparsity <= 1.0 + assert rows % blocking == 0 and cols % blocking == 0 + + block_rows, block_cols = (rows // blocking, cols // blocking) + nnz = round(block_rows * block_cols * (1 - sparsity)) + + out = np.ones(block_rows * block_cols) + mask = np.random.choice(out.size, out.size - nnz, replace=False) + out[mask] = 0.0 + + out = np.tile( + np.reshape(out, [block_rows, 1, block_cols, 1]), + (1, blocking, 1, blocking)) + out = np.reshape(out, [rows, cols]) + return torch.from_numpy(out.astype(np.float32)) + + +@torch.no_grad() +def mask(m, n, sparsity, blocking=1): + out = dense_mask(m, n, sparsity, blocking).type(torch.float16) + return matrix_ops.to_sparse(out, blocking=blocking) + + +@torch.no_grad() +def randn(shape, sparsity, blocking=1): + shape_2d = (np.prod(shape[:-1]), shape[-1]) + out = mask(*shape_2d, sparsity, blocking) + out.data.copy_(torch.randn(*out.data.shape)) + return out.view(*shape) diff --git a/build/torch214-cxx11-cu132-x86_64-linux/stk/random/random_ops_test.py b/build/torch214-cxx11-cu132-x86_64-linux/stk/random/random_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..71d716b78b5ec009cbf9ac2dfdf09162a0102e62 --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/stk/random/random_ops_test.py @@ -0,0 +1,73 @@ +import unittest + +# from absl.testing import parameterized +from . import random +import torch + + +# @parameterized.parameters( +# (8, 16, 0.0, 1), +# (8, 16, 0.5, 1), +# (8, 16, .95, 1), +# (16, 8, 0.0, 1), +# (16, 8, 0.5, 1), +# (16, 8, .95, 1), +# (8, 16, 0.0, 8), +# (8, 16, 0.5, 8), +# (8, 16, 1.0, 8), +# (16, 8, 0.0, 8), +# (16, 8, 0.5, 8), +# (16, 8, 1.0, 8), +# (128, 256, 0.5, 16), +# (256, 128, 0.75, 32), +# (512, 512, .875, 128)) +# class RandomOpsTest(parameterized.TestCase): +# +# def testRandomOps_DenseMask(self, rows, cols, sparsity, blocking): +# mask = random.dense_mask( +# rows, cols, sparsity, blocking) +# +# # Validate the shape. +# self.assertEqual(mask.dim(), 2) +# self.assertEqual(mask.size()[0], rows) +# self.assertEqual(mask.size()[1], cols) +# +# # Validate the sparsity +# numblocks = rows // blocking * cols // blocking +# nnz = round(numblocks * (1 - sparsity)) * blocking ** 2 +# self.assertEqual( +# torch.count_nonzero(mask).item(), +# nnz) +# +# # Check values are zero or one. +# self.assertTrue( +# torch.all(torch.logical_or( +# torch.eq(mask, 0), +# torch.eq(mask, 1)))) +# +# def testRandomOps_SparseMask(self, rows, cols, sparsity, blocking): +# mask = random.mask( +# rows, cols, sparsity, blocking) +# +# # Validate the matrix. +# mask.validate() +# +# # Validate the shape. +# self.assertEqual(mask.dim(), 2) +# self.assertEqual(mask.size()[0], rows) +# self.assertEqual(mask.size()[1], cols) +# +# # Validate the sparsity. +# numblocks = rows // blocking * cols // blocking +# nnz = round(numblocks * (1 - sparsity)) * blocking ** 2 +# self.assertEqual(mask.nnz, nnz) +# +# # Check values are zero or one. +# self.assertTrue( +# torch.all(torch.logical_or( +# torch.eq(mask.data, 0), +# torch.eq(mask.data, 1)))) + + +if __name__ == '__main__': + unittest.main() diff --git a/build/torch214-cxx11-cu132-x86_64-linux/xpu_fused_moe.py b/build/torch214-cxx11-cu132-x86_64-linux/xpu_fused_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..b709e394e2e0461d858dec0f58176c0543a15e3c --- /dev/null +++ b/build/torch214-cxx11-cu132-x86_64-linux/xpu_fused_moe.py @@ -0,0 +1,717 @@ +# SPDX-License-Identifier: Apache-2.0 +# MegaBlocks XPU Fused MoE Implementation +import os +import torch + +from ._ops import ops, add_op_namespace_prefix + +from torch.library import register_fake + + +def resolve_dtensor(weight: torch.Tensor): + """Convert DTensor to local tensor for use with custom ops.""" + from torch.distributed._tensor import DTensor + if isinstance(weight, DTensor): + return weight.to_local() + return weight + + +# Register fake/meta kernels for torch.compile compatibility +def _register_xpu_fake_kernels(): + """Register fake kernels for XPU MoE operations to support torch.compile.""" + + def _register_if_available(op_name, fn): + if hasattr(ops, op_name): + register_fake(add_op_namespace_prefix(op_name))(fn) + + _register_if_available( + "cutlass_grouped_gemm_interface", + lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4, is_B_mxfp8=False: ptr_D, + ) + + _register_if_available( + "fused_moe_prologue", + lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, ep_rank, ep_size, num_experts_on_rank: None, + ) + + _register_if_available( + "moe_gather", + lambda output, moe_output, topk_weights, permuted_row_to_unpermuted_row, unpermuted_row_to_permuted_row, expert_first_token_offset, num_experts: None, + ) + + _register_if_available( + "silu_and_mul", + lambda out, input: None, + ) + _register_if_available( + "mul_and_silu", + lambda out, input: None, + ) + _register_if_available( + "gelu_and_mul", + lambda out, input: None, + ) + _register_if_available( + "gelu_tanh_and_mul", + lambda out, input: None, + ) + _register_if_available( + "gelu_fast", + lambda out, input: None, + ) + _register_if_available( + "gelu_new", + lambda out, input: None, + ) + _register_if_available( + "gelu_quick", + lambda out, input: None, + ) + _register_if_available( + "swigluoai_and_mul", + lambda out, input, alpha=1.702, limit=7.0: None, + ) + + +# Register fake kernels on module load +_register_xpu_fake_kernels() + + +# default +def cutlass_grouped_gemm(input_A, input_B, bias, output, expert_token_count, n, + k, num_experts): + # expert_token_count_ = torch.tensor(expert_token_count, + # dtype=torch.int64, + # device=input_A.device) + # if bias is not None: + # bias = bias.repeat_interleave(expert_token_count_, dim=0).float() + + def exclusive_prefix_sum(arr): + prefix = [0] + for i, x in enumerate(arr): + prefix.append(prefix[-1] + x) + return prefix + + expert_offset = torch.tensor(exclusive_prefix_sum(expert_token_count), + dtype=torch.int64, + device="xpu") + ops.cutlass_grouped_gemm_interface( + ptr_A=input_A, + ptr_B=input_B, + ptr_scales=None, + ptr_bias=bias, + ptr_D=output, + expert_first_token_offset=expert_offset, + N=n, + K=k, + num_experts=num_experts, + is_B_int4=False, + is_B_mxfp4=False, + is_B_mxfp8=False) + + +def cutlass_grouped_gemm_xe2(input_A, input_B, scales, bias, output, + num_rows_per_expert, n, k, num_experts, is_B_int4, + is_B_mxfp4, is_B_mxfp8=False): + expert_first_token_offset = torch.cat([ + torch.tensor([0], + dtype=num_rows_per_expert.dtype, + device=num_rows_per_expert.device), + torch.cumsum(num_rows_per_expert, dim=0) + ]).to(torch.int64) + ops.cutlass_grouped_gemm_interface( + ptr_A=input_A, + ptr_B=input_B, + ptr_scales=scales, + ptr_bias=bias, + ptr_D=output, + expert_first_token_offset=expert_first_token_offset, + N=n, + K=k, + num_experts=num_experts, + is_B_int4=is_B_int4, + is_B_mxfp4=is_B_mxfp4, + is_B_mxfp8=is_B_mxfp8) + + +def ceilDiv(a, b): + return (a + b - 1) // b + + +def compute_num_tokens_per_block(num_tokens, num_experts_per_node): + for num_tokens_per_block in [32, 64, 128, 256, 512, 1024]: + num_blocks_per_seq = ceilDiv(num_tokens, num_tokens_per_block) + if num_blocks_per_seq * num_experts_per_node <= num_tokens_per_block: + return num_tokens_per_block + return 1024 + + +def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """Reinterpret a uint8 buffer as a typed tensor by copying bytes. + + This avoids `Tensor.view(dtype)` which can fail under torch.compile + constant folding when shape divisibility is not proven. + """ + if byte_tensor.dtype != torch.uint8: + raise ValueError("byte_tensor must be uint8") + itemsize = torch.empty((), dtype=dtype).element_size() + numel = byte_tensor.numel() // itemsize + out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device) + out.view(torch.uint8).copy_(byte_tensor.contiguous()) + return out + + +def implement_zp(qweight): + # change u4 to s4 to avoid zero point in gemm kernel + # only support default zero point now + assert qweight.dtype == torch.uint8, "Input tensor must be uint8" + + high_u4 = (qweight >> 4) & 0x0F + low_u4 = qweight & 0x0F + + high_s8 = high_u4.to(torch.int8) + low_s8 = low_u4.to(torch.int8) + + high_s8 = high_s8 - 8 + low_s8 = low_s8 - 8 + + def pack_compact(a, b): + + def process_number(x): + sign = (x < 0).to(torch.uint8) + abs_low3 = (x.view(torch.uint8) & 0x7).to(torch.uint8) + return (sign << 3) | abs_low3 + + packed_a = process_number(a) + packed_b = process_number(b) + + return (packed_a << 4) | packed_b + + result = pack_compact(high_s8, low_s8) + + return result + + +def xpu_fused_moe(hidden_states, + w13, + w13_scales, + w13_bias, + w2, + w2_scales, + w2_bias, + topk_weights, + topk_ids, + n_experts_per_token, + activation, + num_experts, + ep_rank=0, + ep_size=1, + is_fp8=False, + is_int4=False, + is_mxfp4=False, + is_mxfp8=False): + ''' + hidden_states: [num_rows, hidden_size] + w13: [num_experts, 2*inter_size, hidden_size] + w13_scales: + None for bf16/fp16 + or [num_experts] for fp8 + or [num_experts, 2*inter_size, hidden_size // group_size] for 4bits + w13_bias: [num_experts, 2*inter_size] or None + w2: [num_experts, hidden_size, inter_size] + w2_scales: + None for bf16/fp16 + or [num_experts] for fp8 + or [num_experts, hidden_size, inter_size // group_size] for 4bits + w2_bias: [num_experts, hidden_size] or None + topk_weights: [num_rows, topk] + topk_ids: [num_rows, topk] + n_experts_per_token: int + activation: str + num_experts: int + is_int4: bool + is_mxfp4: bool + is_mxfp8: bool + ''' + + # Resolve DTensors to local tensors before passing to custom ops + hidden_states = resolve_dtensor(hidden_states) + w13 = resolve_dtensor(w13) + w2 = resolve_dtensor(w2) + if w13_scales is not None: + w13_scales = resolve_dtensor(w13_scales) + if w13_bias is not None: + w13_bias = resolve_dtensor(w13_bias) + if w2_scales is not None: + w2_scales = resolve_dtensor(w2_scales) + if w2_bias is not None: + w2_bias = resolve_dtensor(w2_bias) + topk_weights = resolve_dtensor(topk_weights) + topk_ids = resolve_dtensor(topk_ids) + + output = torch.empty_like(hidden_states) + num_rows, hidden_size = list(hidden_states.shape) + + dim_last = w13.shape[-1] + dim_second_last = w13.shape[-2] + + # w13 is combined gate+up weights, so one dimension is 2*inter_size + # Determine which dimension is hidden_size and which is 2*inter_size + if dim_second_last == hidden_size: + # w13 is [E, hidden_size, 2*inter_size] - standard layout + inter_size = dim_last // 2 + needs_transpose = False + else: + # w13 is [E, 2*inter_size, hidden_size] - needs transpose + inter_size = dim_second_last // 2 + needs_transpose = True + + assert w13.is_contiguous() and w2.is_contiguous() + + # 4bits support [E, N, K] + # other types [E, K, N] + if not is_int4 and not is_mxfp4 and not is_mxfp8: + if not hasattr(w13, 'xpu_fused_moe'): + if needs_transpose: + w13.data = w13.transpose(-1, -2).contiguous() + w2.data = w2.transpose(-1, -2).contiguous() + w13.xpu_fused_moe = True + w13.inter_size = inter_size + else: + inter_size = w13.inter_size + + if is_int4 and not hasattr(w13, 'xpu_fused_moe'): + w13_tmp = torch.empty_like(w13) + w2_tmp = torch.empty_like(w2) + for i in range(num_experts): + w13_tmp[i] = implement_zp(w13[i]) + w2_tmp[i] = implement_zp(w2[i]) + w13_tmp = w13_tmp.contiguous() + w2_tmp = w2_tmp.contiguous() + w13.data = w13_tmp + w2.data = w2_tmp + w13.xpu_fused_moe = True + + # TODO: will all integrated in Cpp func. Temporary expose before gemm fusion + num_experts_per_node = num_experts + experts_per_token = n_experts_per_token + num_moe_inputs = n_experts_per_token * num_rows + permuted_elems = num_moe_inputs * hidden_size + # interbuf_elems = num_moe_inputs * inter_size + permuted_row_to_unpermuted_row_size = num_moe_inputs * 4 + permuted_token_selected_experts_size = num_moe_inputs * 4 + src_to_dest_map_size = experts_per_token * num_rows * 4 + expert_first_token_offset_size = (num_experts_per_node + 1) * 8 + num_tokens_per_block = compute_num_tokens_per_block( + num_rows, num_experts_per_node) + num_blocks_per_seq = ceilDiv(num_rows, num_tokens_per_block) + blocked_expert_counts_size = num_experts_per_node * num_blocks_per_seq * 4 + blocked_expert_counts_cumsum_size = blocked_expert_counts_size + blocked_row_to_unpermuted_row_size = num_experts_per_node * num_rows * 4 + permuted_data_size = permuted_elems * hidden_states.element_size() + permuted_token_final_scales_size = num_moe_inputs * 4 + + ws_map = {} + map_offset = 0 + + def config_ws(name, size): + nonlocal map_offset + if size % 256 != 0: + size += 256 - size % 256 + ws_map[name] = (size, map_offset) + map_offset += size + + config_ws("permuted_row_to_unpermuted_row", + permuted_row_to_unpermuted_row_size) + config_ws("permuted_token_selected_experts", + permuted_token_selected_experts_size) + config_ws("unpermuted_row_to_permuted_row", src_to_dest_map_size) + config_ws("blocked_expert_counts", blocked_expert_counts_size) + config_ws("blocked_expert_counts_cumsum", + blocked_expert_counts_cumsum_size) + config_ws("blocked_row_to_unpermuted_row", + blocked_row_to_unpermuted_row_size) + config_ws("expert_first_token_offset", expert_first_token_offset_size) + config_ws("permuted_token_final_scales", permuted_token_final_scales_size) + config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size) + + workspace = torch.zeros(map_offset, + dtype=torch.uint8, + device=hidden_states.device) + if topk_ids.dtype == torch.int32: + topk_ids = topk_ids.to(torch.int64) + ops.fused_moe_prologue( + input=hidden_states, + token_selected_experts=topk_ids, + token_final_scales=topk_weights, + workspace=workspace, + hidden_size=hidden_size, + inter_size=inter_size, + ep_rank=ep_rank, + ep_size=ep_size, + num_experts_on_rank=num_experts_per_node) + + expert_first_token_offset_bytes = workspace[ + ws_map["expert_first_token_offset"][1]: + ws_map["expert_first_token_offset"][1] + + expert_first_token_offset_size] + unpermuted_row_to_permuted_row_bytes = workspace[ + ws_map["unpermuted_row_to_permuted_row"][1]: + ws_map["unpermuted_row_to_permuted_row"][1] + + src_to_dest_map_size] + permuted_row_to_unpermuted_row_bytes = workspace[ + ws_map["permuted_row_to_unpermuted_row"][1]: + ws_map["permuted_row_to_unpermuted_row"][1] + + permuted_row_to_unpermuted_row_size] + + if torch.compiler.is_compiling(): + expert_first_token_offset = _bytes_to_typed_tensor( + expert_first_token_offset_bytes, torch.int64 + ) + unpermuted_row_to_permuted_row = _bytes_to_typed_tensor( + unpermuted_row_to_permuted_row_bytes, torch.int32 + ) + permuted_row_to_unpermuted_row = _bytes_to_typed_tensor( + permuted_row_to_unpermuted_row_bytes, torch.int32 + ) + else: + expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64) + unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32) + permuted_row_to_unpermuted_row = permuted_row_to_unpermuted_row_bytes.view(torch.int32) + gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]: + ws_map["overlapped_gemm1_gemm2_inputs"][1] + + permuted_data_size].view(hidden_states.dtype).view( + num_moe_inputs, hidden_size) + # permuted_token_final_scales = workspace[ + # ws_map["permuted_token_final_scales"][1]: + # ws_map["permuted_token_final_scales"][1] + + # permuted_token_final_scales_size].view(torch.float) + gemm1_output = torch.empty((num_moe_inputs, 2 * inter_size), + dtype=hidden_states.dtype, + device=hidden_states.device) + + ########### gemm1 ################## + input_B = w13 + + if not is_fp8 and not is_int4 and not is_mxfp4 and not is_mxfp8: + ops.cutlass_grouped_gemm_interface( + ptr_A=gemm1_input, + ptr_B=input_B, + ptr_scales=None, + ptr_bias=w13_bias, + ptr_D=gemm1_output, + expert_first_token_offset=expert_first_token_offset, + N=2 * inter_size, + K=hidden_size, + num_experts=num_experts_per_node, + is_B_int4=is_int4, + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) + else: + ops.cutlass_grouped_gemm_interface( + ptr_A=gemm1_input, + ptr_B=input_B, + ptr_scales=w13_scales, + ptr_bias=w13_bias, + ptr_D=gemm1_output, + expert_first_token_offset=expert_first_token_offset, + N=2 * inter_size, + K=hidden_size, + num_experts=num_experts_per_node, + is_B_int4=is_int4, + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) + + # act + act_output = torch.empty((num_moe_inputs, inter_size), + dtype=gemm1_output.dtype, + device=gemm1_output.device) + if activation == "silu": + ops.silu_and_mul(act_output, gemm1_output) + elif activation == "gelu": + ops.gelu_and_mul(act_output, gemm1_output) + elif activation == "swigluoai": + ops.swigluoai_and_mul(act_output, gemm1_output, 1.702, 7.0) + else: + raise ValueError(f"Unsupported FusedMoe activation: {activation}.") + + ########### gemm2 ################## + input_A = act_output.contiguous() + input_B = w2 + gemm2_output = torch.empty((num_moe_inputs, hidden_size), + dtype=hidden_states.dtype, + device=hidden_states.device) + if not is_fp8 and not is_int4 and not is_mxfp4 and not is_mxfp8: + ops.cutlass_grouped_gemm_interface( + ptr_A=input_A, + ptr_B=input_B, + ptr_scales=None, + ptr_bias=w2_bias, + ptr_D=gemm2_output, + expert_first_token_offset=expert_first_token_offset, + N=hidden_size, + K=inter_size, + num_experts=num_experts_per_node, + is_B_int4=is_int4, + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) + else: + ops.cutlass_grouped_gemm_interface( + ptr_A=input_A, + ptr_B=input_B, + ptr_scales=w2_scales, + ptr_bias=w2_bias, + ptr_D=gemm2_output, + expert_first_token_offset=expert_first_token_offset, + N=hidden_size, + K=inter_size, + num_experts=num_experts_per_node, + is_B_int4=is_int4, + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) + + ops.moe_gather(output, gemm2_output, topk_weights, + permuted_row_to_unpermuted_row, + unpermuted_row_to_permuted_row, + expert_first_token_offset, + num_experts_per_node) + return output + + +def apply_jitter(x: torch.Tensor, moe_jitter_eps: float) -> torch.Tensor: + """Apply jitter to the input tensor for regularization.""" + low = 1.0 - moe_jitter_eps + high = 1.0 + moe_jitter_eps + noise = torch.rand(x.size(), dtype=x.dtype, device=x.device) + return x * (low + noise * (high - low)) + + +def compute_top_k(scores: torch.Tensor, moe_top_k: int): + """Compute the top-k scores from the logits.""" + if moe_top_k == 1: + return scores.max(dim=-1, keepdim=True) + return torch.topk(scores, moe_top_k, dim=-1) + + +def route_tokens_xpu( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + moe_top_k: int, + moe_num_experts: int, + moe_jitter_eps: float = None, + moe_normalize_expert_weights: int = None, + training: bool = False, +) -> tuple: + """Route tokens to experts and compute expert weights and indices (XPU version).""" + if training and moe_jitter_eps is not None: + x = apply_jitter(x, moe_jitter_eps) + + x_flat = x.view(-1, x.shape[-1]) + logits = torch.nn.functional.linear(x_flat, router_weight, router_bias) + expert_weights, expert_indices = compute_top_k(logits, moe_top_k) + expert_weights = expert_weights.softmax(dim=-1) + if moe_normalize_expert_weights is not None: + expert_weights = expert_weights / torch.norm( + expert_weights, + p=moe_normalize_expert_weights, + dim=-1, + keepdim=True, + ) + + return logits, expert_weights, expert_indices + + +def _get_device_mesh(model): + """Extract device_mesh from child's unused pre_hook closure for EP support.""" + try: + hook = next( + h + for h in model.experts._forward_pre_hooks.values() + if "device_mesh" in h.__code__.co_freevars + ) + return hook.__closure__[ + hook.__code__.co_freevars.index("device_mesh") + ].cell_contents + except Exception: + return None + + +def _mxfp4_expert_tensors(experts, dtype): + """View transformers' MXFP4 experts as the tensors the SYCL grouped GEMM expects. + + Weights and scales arrive wrapped in triton tensors. On XPU both use StridedLayout, + whose swizzle is the identity, so `storage.data` is still the checkpoint tensor - just + transposed to [E, K/2, N]. Transposing back therefore restores the original contiguous + [E, N, K/2] weights and [E, N, K/32] E8M0 scales as plain views: `contiguous()` finds + them already contiguous and returns them untouched, so this costs no copy and no extra + memory. The experts module is left alone, which keeps the triton tensors and the + PrecisionConfig that `save_pretrained` needs to re-serialize the checkpoint. + """ + tensors = [] + for proj in ("gate_up_proj", "down_proj"): + config = getattr(experts, f"{proj}_precision_config") + tensors.append(getattr(experts, proj).storage.data.transpose(-1, -2).contiguous()) + tensors.append(config.weight_scale.storage.data.transpose(-1, -2).contiguous()) + + # The GEMM reads the bias as ElementA; `to` is a no-op when it already matches. + bias = getattr(experts, f"{proj}_bias", None) + tensors.append(None if bias is None else bias.data.to(dtype)) + + return tensors + + +class MegaBlocksMoeMLP(torch.nn.Module): + can_torch_compile: bool = True + + def forward(self, x: torch.Tensor) -> tuple: + """ + Forward pass through the MoE layer. + + Args: + x: Input tensor of shape [batch_size, seq_len, hidden_size] or [tokens, hidden_size] + + Returns: + Tuple of (output, expert_weights) where: + - output: Tensor of same shape as input + - expert_weights: Expert weights for each token [tokens, top_k] + """ + + # Get MoE parameters from the wrapped modules + moe_top_k = getattr(self.router, "top_k", 4) + moe_num_experts = getattr(self.experts, "num_experts", 128) + moe_jitter_eps = getattr(self.experts, "jitter_eps", None) + moe_normalize_expert_weights = getattr( + self.experts, "normalize_expert_weights", None + ) + + # Get EP (Expert Parallelism) parameters + ep_size = 1 + ep_rank = 0 + expert_parallel_group = getattr(self, "expert_parallel_group", None) + if expert_parallel_group is None: + device_mesh = _get_device_mesh(self) + if device_mesh is not None: + expert_parallel_group = device_mesh.get_group() + if expert_parallel_group is not None: + import torch.distributed as dist + if dist.is_initialized(): + ep_size = dist.get_world_size(expert_parallel_group) + ep_rank = dist.get_rank(expert_parallel_group) + + # Number of experts on this rank + num_experts_on_rank = moe_num_experts // ep_size + + # Detect activation type - check for GptOss-style swigluoai activation + # GptOssExperts has alpha and limit attributes for swigluoai + if hasattr(self.experts, "alpha") and hasattr(self.experts, "limit"): + activation = "swigluoai" + else: + activation = getattr(self.experts, "activation", "silu") + + # Get weight tensors - support different naming conventions + if hasattr(self.experts, "gate_up_proj"): + w13 = self.experts.gate_up_proj + # NOTE: swigluoai_and_mul kernel expects interleaved layout [g0,u0,g1,u1,...] + # which matches GptOss's gate_up_proj layout, so no conversion needed. + + elif hasattr(self.experts, "w1"): + # Combine w1 and w3 if stored separately + w1 = self.experts.w1 + w3 = getattr(self.experts, "w3", None) + if w3 is not None: + w13 = torch.cat([w1, w3], dim=-2) + else: + w13 = w1 + else: + raise AttributeError("experts module must have 'gate_up_proj' or 'w1' attribute") + + if hasattr(self.experts, "down_proj"): + w2 = self.experts.down_proj + elif hasattr(self.experts, "w2"): + w2 = self.experts.w2 + else: + raise AttributeError("experts module must have 'down_proj' or 'w2' attribute") + + # Get optional bias tensors + w13_bias = getattr(self.experts, "gate_up_proj_bias", None) + w2_bias = getattr(self.experts, "down_proj_bias", None) + + # Get quantization info + is_fp8 = getattr(self.experts, "is_fp8", False) + is_int4 = getattr(self.experts, "is_int4", False) + is_mxfp4 = getattr(self.experts, "is_mxfp4", False) + is_mxfp8 = getattr(self.experts, "is_mxfp8", False) + + w13_scales = getattr(self.experts, "gate_up_proj_scales", None) + w2_scales = getattr(self.experts, "down_proj_scales", None) + + # Native MXFP4 checkpoint: derive the GEMM's view of the experts once and cache it. + mxfp4 = getattr(self, "_mxfp4_tensors", None) + if (mxfp4 is None or mxfp4[0] is not x.dtype) and getattr( + self.experts, "gate_up_proj_precision_config", None + ) is not None: + mxfp4 = (x.dtype, *_mxfp4_expert_tensors(self.experts, x.dtype)) + self._mxfp4_tensors = mxfp4 + if mxfp4 is not None: + _, w13, w13_scales, w13_bias, w2, w2_scales, w2_bias = mxfp4 + is_mxfp4 = True + + # Store original shape + in_shape = x.size() + + # Route tokens to experts + logits, expert_weights, expert_indices = route_tokens_xpu( + x, + self.router.weight, + getattr(self.router, "bias", None), + moe_top_k, + moe_num_experts, + moe_jitter_eps, + moe_normalize_expert_weights, + self.training, + ) + + # Reshape input for fused MoE + x_flat = x.view(-1, x.shape[-1]) + + # Call XPU fused MoE kernel + output = xpu_fused_moe( + hidden_states=x_flat, + w13=w13, + w13_scales=w13_scales, + w13_bias=w13_bias, + w2=w2, + w2_scales=w2_scales, + w2_bias=w2_bias, + topk_weights=expert_weights.float(), + topk_ids=expert_indices, + n_experts_per_token=moe_top_k, + activation=activation, + num_experts=num_experts_on_rank, + ep_rank=ep_rank, + ep_size=ep_size, + is_fp8=is_fp8, + is_int4=is_int4, + is_mxfp4=is_mxfp4, + is_mxfp8=is_mxfp8, + ) + + # All-reduce across EP group to combine partial expert outputs + if ep_size > 1 and expert_parallel_group is not None: + import torch.distributed as dist + dist.all_reduce(output, op=dist.ReduceOp.SUM, group=expert_parallel_group) + + # Restore original shape + output = output.view(in_shape) + + return output, expert_weights + + +# Export classes and functions +__all__ = [ + "MegaBlocksMoeMLP", + "xpu_fused_moe", + "cutlass_grouped_gemm", + "cutlass_grouped_gemm_xe2", +] \ No newline at end of file