Kernels:

kernels-community
/

megablocks

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on Aug 27

Commit

2c977e2

verified ·

1 Parent(s): 65478c1

Uploaded using `kernel-builder`.

Browse files

This view is limited to 50 files because it contains too many changes. See raw diff

Files changed (50) hide show

build/torch212-cxx11-xpu20253-x86_64-linux/{_megablocks_xpu_c4d0dc5.abi3.so → _megablocks_xpu_addf474.abi3.so} +1 -1
build/torch212-cxx11-xpu20253-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-xpu20253-x86_64-linux/megablocks/__init__.py +0 -26
build/torch212-cxx11-xpu20253-x86_64-linux/metadata.json +6 -6
build/torch212-cxx11-xpu20253-x86_64-linux/metadata.json.sigstore +1 -1
build/torch213-cxx11-xpu20260-x86_64-linux/{_megablocks_xpu_c4d0dc5.abi3.so → _megablocks_xpu_addf474.abi3.so} +2 -2
build/torch213-cxx11-xpu20260-x86_64-linux/_ops.py +3 -3
build/torch213-cxx11-xpu20260-x86_64-linux/megablocks/__init__.py +0 -26
build/torch213-cxx11-xpu20260-x86_64-linux/metadata.json +6 -6
build/torch213-cxx11-xpu20260-x86_64-linux/metadata.json.sigstore +1 -1
build/torch214-cxx11-xpu20261-x86_64-linux/__init__.py +205 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/__init__.py +0 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/_triton_kernels/__init__.py +0 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/_triton_kernels/gmm.py +574 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/adapter.py +53 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/configs.py +5 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/gmm.py +567 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/utils/__init__.py +0 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/utils/_triton/__init__.py +0 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/utils/_triton/arch_info.py +46 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/utils/_triton/pid_preprocessing.py +100 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/utils/gmm_common.py +752 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/utils/logger.py +47 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/__init__.py +10 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/activation_fn.py +33 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/all_to_all.py +54 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/arguments.py +101 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/common.py +26 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/dmlp_registry.py +42 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/dmoe.py +337 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/gelu.py +52 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/glu.py +244 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/memory_test.py +103 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/mlp.py +587 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/moe.py +507 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/mpu.py +94 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/router.py +116 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_layers/sharedexpert_registry.py +32 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_megablocks_xpu_addf474.abi3.so +3 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_ops.py +9 -0
build/torch214-cxx11-xpu20261-x86_64-linux/_version.py +6 -0
build/torch214-cxx11-xpu20261-x86_64-linux/backend/__init__.py +2 -0
build/torch214-cxx11-xpu20261-x86_64-linux/backend/kernels.py +557 -0
build/torch214-cxx11-xpu20261-x86_64-linux/benchmark_util.py +35 -0
build/torch214-cxx11-xpu20261-x86_64-linux/cpu_fused_moe.py +311 -0
build/torch214-cxx11-xpu20261-x86_64-linux/cpu_moe_cpp.py +265 -0
build/torch214-cxx11-xpu20261-x86_64-linux/grouped_gemm/__init__.py +2 -0
build/torch214-cxx11-xpu20261-x86_64-linux/grouped_gemm/backend.py +33 -0
build/torch214-cxx11-xpu20261-x86_64-linux/grouped_gemm/ops.py +33 -0
build/torch214-cxx11-xpu20261-x86_64-linux/grouped_gemm_util.py +31 -0

build/torch212-cxx11-xpu20253-x86_64-linux/{_megablocks_xpu_c4d0dc5.abi3.so → _megablocks_xpu_addf474.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:868eb695bd1e4c252332321da1ca3002e7322764498808ed957cf8e869b9ea5a
 size 4321128

 version https://git-lfs.github.com/spec/v1
+oid sha256:69446c0ed6ada667cbc27e2cf08f257dd2d3e9955657d5ac87eff09d3350ed3d
 size 4321128

build/torch212-cxx11-xpu20253-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_xpu_c4d0dc5
-ops = torch.ops._megablocks_xpu_c4d0dc5
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_xpu_c4d0dc5::{op_name}"

 import torch
+from . import _megablocks_xpu_addf474
+ops = torch.ops._megablocks_xpu_addf474
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_megablocks_xpu_addf474::{op_name}"

build/torch212-cxx11-xpu20253-x86_64-linux/megablocks/__init__.py DELETED Viewed

@@ -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")))

build/torch212-cxx11-xpu20253-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,9 +1,10 @@
 {
   "name": "megablocks",
-  "id": "_megablocks_xpu_c4d0dc5",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "xpu"
   },
@@ -38,8 +39,8 @@
       "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=",
       "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=",
       "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=",
-      "_megablocks_xpu_c4d0dc5.abi3.so": "ho62lb0eTCUjMjIdocowAucyJ2RJiAjtlXz46Gm56lo=",
-      "_ops.py": "KZ/xutl7d4r1cvlnq+lt8PGsR1a57OQznlwnP3eVkKI=",
       "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=",
       "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=",
       "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=",
@@ -51,7 +52,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=",
@@ -96,11 +96,11 @@
   "provenance": {
     "kernel-builder": {
       "version": "0.17.0-dev0",
-      "sha": "81580bb92577f2f7228661ca2a221fc052375709",
       "dirty": false
     },
     "kernel": {
-      "sha": "c4d0dc56661badb9cc47bde647e08fe3b4bf7458",
       "dirty": false
     }
   }

 {
   "name": "megablocks",
+  "id": "_megablocks_xpu_addf474",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
+  "kernel-depends": [],
   "backend": {
     "type": "xpu"
   },
       "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=",
       "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=",
       "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=",
+      "_megablocks_xpu_addf474.abi3.so": "aURsDtatpmfLwn4s8I8lfdLT6ZVWV9Wsh+/wnTNQ7T0=",
+      "_ops.py": "QBFLHcpmqB9+JkTkhdpkiEE9yER4FUDrsZgGzybNoVc=",
       "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=",
       "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=",
       "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=",
       "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=",
   "provenance": {
     "kernel-builder": {
       "version": "0.17.0-dev0",
+      "commit": "7d2828e6592de641b44fb8eb896719101a2ab2fb",
       "dirty": false
     },
     "kernel": {
+      "commit": "addf4741fafb55b4c101b15985dfaeea7c7c2f3a",
       "dirty": false
     }
   }

build/torch212-cxx11-xpu20253-x86_64-linux/metadata.json.sigstore CHANGED Viewed

@@ -1 +1 @@

- {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHTDCCBtKgAwIBAgIUXl5Lj0ZZxmJSBAElnUur5bPOFSUwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI1MTIwNzQ3WhcNMjYwODI1MTIxNzQ3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAErc2vjnOd2F7gDDhqBurUUxT/ujdwHbHjt8JcM57UgNubELzPeSCPV+EjHNgld7qmRfUjLa21P6DWD7BypZNRNqOCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU+Uxv/tFJgmXCXe5vuYfg2ML3TocwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGM0ZDBkYzU2NjYxYmFkYjljYzQ3YmRlNjQ3ZTA4ZmUzYjRiZjc0NTgwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzI4NDQ2ODk4OTQvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoDjRxswAAAQDAEgwRgIhAN+W8o+X7UgM1dt/DI3YoGQwFZtyRfMAGOEyZnvXz5DPAiEA0zBGzcriJLXqkz+htQgXqRdIzS3LI70Vj61poOh2B+MwCgYIKoZIzj0EAwMDaAAwZQIxANO7WbSfiTPHzqSEAfB+JpfkSWRdxfujCm2jE0ncT/4jIC6v14qyZVBlHp7+teGYYAIwOx4kJzvP6rP/tlp5LPhe7rF5kD+Vja6AtjwfMNXLBgTcU4Aew0pof22Kdj5DsGO9"}, "tlogEntries":[{"logIndex":"2583274380", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787659667", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQCpPbM0AS/szLFi2f2RNpUPV4OvvmCRi5a6WAzN9iI8/AIhAJA2d6b7uUNN4ebditgvyIz/6tvKd7yv4hi/AP2Mwxti"}, "inclusionProof":{"logIndex":"2461370118", "rootHash":"cQ7wW1xW/L20rkOvTt9r9iCKPByWXNI6FCkRaju37Ts=", "treeSize":"2461370120", "hashes":["y7dnw8OA1nxLA6kcdGFSV0dFZoGA1uZ1xOITOKTf2Js=", "uJ0tsePLf6AS/EahjuJN0+b+EnlFud7NirVAZMI0B4I=", "jjdaExfTfO3Htve6MMlB8UBrmhQIgQW9lSQ4kVJYuG4=", "BJlxWyo+0U2jxgcimrAeK1pey6FSpHWLOetp0UXRWI4=", "JsX9xGIEwFv7MFelTobmTLz4j1Jl78aSs2G7dkybw0w=", "nwobU/9yRWhsbw80xq3x/uNV1nniMWdUKAqc3UWYsJM=", "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\n2461370120\ncQ7wW1xW/L20rkOvTt9r9iCKPByWXNI6FCkRaju37Ts=\n\n— rekor.sigstore.dev wNI9ajBFAiA9mr274M9R9qIT6J8dTEeSqj88FaQxjItF4bRLmWCsmwIhAMsWTJ1A8HcsmgMoEJo+YpOQ01H1eQMrBHPvtx3pT2mT\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJlMjU4NGVmMDU2OGZkMmVkMTM4MmE2ODZjODNhMGIyNzNhMmZjYzE4MWI1YzZkMmQ2ODRjNzgwNDc5MzNjM2JkIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNwRlg3N3NjWTF1VUZwOWdNQWIrUWJ3VVluYlNFdXNjYURBYlJCbFVQNHJnSWdCT1E3T3U4QnNaTGo1RWNtNDArOTViOXZoUFhGVmZZNmxxV3kvUHpBbWE0PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVSRU5EUW5STFowRjNTVUpCWjBsVldHdzFUR293V2xwNGJVcFRRa0ZGYkc1VmRYSTFZbEJQUmxOVmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1RGTlZFbDNUbnBSTTFkb1kwNU5hbGwzVDBSSk1VMVVTWGhPZWxFelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ5WXpKMmFtNVBaREpHTjJkRVJHaHhRblZ5VlZWNFZDOTFhbVIzU0dKSWFuUTRTbU1LVFRVM1ZXZE9kV0pGVEhwUVpWTkRVRllyUldwSVRtZHNaRGR4YlZKbVZXcE1ZVEl4VURaRVYwUTNRbmx3V2s1U1RuRlBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlVyVlhoMkNpOTBSa3BuYlZoRFdHVTFkblZaWm1jeVRVd3pWRzlqZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVplbEpyVFVkU2FrNVVXVEpPYWtacFdWZFNhVTlYVG1wT1JHUnBDbHBIVlRKT1JHUnNUVVJvYlZwVVRtbE9SMHB0VG5wUk1VOUVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2xhUjFVeVRrUmtiRTFFYUcwS1dsUk9hVTVIU20xT2VsRXhUMFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2tLV2tkVk1rNUVaR3hOUkdodFdsUk9hVTVIU20xT2VsRXhUMFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZE5NRnBFUW1zS1dYcFZNazVxV1hoWmJVWnJXV3BzYWxsNlVUTlpiVkpzVG1wUk0xcFVRVFJhYlZWNldXcFNhVnBxWXpCT1ZHZDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtrMFRrUlJNazlFYXpSUFZGRjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUVhbEo0YzNkQlFVRlJSQXBCUldkM1VtZEphRUZPSzFjNGJ5dFlOMVZuVFRGa2RDOUVTVE5aYjBkUmQwWmFkSGxTWmsxQlIwOUZlVnB1ZGxoNk5VUlFRV2xGUVRCNlFrZDZZM0pwQ2twTVdIRnJlaXRvZEZGbldIRlNaRWw2VXpOTVNUY3dWbW8yTVhCdlQyZ3lRaXROZDBObldVbExiMXBKZW1vd1JVRjNUVVJoUVVGM1dsRkplRUZPVHpjS1YySlRabWxVVUVoNmNWTkZRV1pDSzBwd1ptdFRWMUprZUdaMWFrTnRNbXBGTUc1alZDODBha2xETm5ZeE5IRjVXbFpDYkVod055dDBaVWRaV1VGSmR3cFBlRFJyU25wMlVEWnlVQzkwYkhBMVRGQm9aVGR5UmpWclJDdFdhbUUyUVhScWQyWk5UbGhNUW1kVVkxVTBRV1YzTUhCdlpqSXlTMlJxTlVSelIwODVDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgAF4NzZ4QDxxyny7GD5G9gutsEM6gS48lZIjdv0q0L2oCFCqbR84O99XA9ORjAz6eaf5ry4FfGA8yMDI2MDgyNTEyMDc0N1owAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwODI1MTIwNzQ3WjAvBgkqhkiG9w0BCQQxIgQg0xV8kuoZmTLPREuzDn9A2nFk6uc+jwxngJv06AITOqwwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAMORtmStWCwKwKsaiLGpF8sh0OG1iC4tOrFRmmNktIbLmI9CIQUsE5ApgIm8LWbMwAIxAOEkQlSqUSkaNMPcj+l3vymorX1kyGFtf1kzOX789cV5zvO8oJs+Z4Ob8ZAs+jdMhQ=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"4lhO8FaP0u0TgqaGyDoLJzovzBgbXG0taEx4BHkzw70="}, "signature":"MEUCIQCpFX77scY1uUFp9gMAb+QbwUYnbSEuscaDAbRBlUP4rgIgBOQ7Ou8BsZLj5Ecm40+95b9vhPXFVfY6lqWy/PzAma4="}}

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHTTCCBtKgAwIBAgIUeFaEsDg9TNOkT8jPh1eH7XW1ZG0wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI3MDcwMTMzWhcNMjYwODI3MDcxMTMzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE3Mql367h5c67VBxvbhUlurjO4WCr8b8hEpLd/9c1GxY5TKhxm/QZb3/NAioZJru7lljD/Lq62har3sLL0gz2J6OCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUdcb1KKYaYn60xVmuNMNRUFJtLsQwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGFkZGY0NzQxZmFmYjU1YjRjMTAxYjE1OTg1ZGZhZWVhN2M3YzJmM2EwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzMwNDUzOTU5NTIvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoEIGI2UAAAQDAEgwRgIhAMEl1OtXrc+rgzO8F3FC5n353ulp7ViQeVkv/zbKyx+zAiEAw3YdtCwMf1v8wCiqHIkKwR8cwG1hKVQ6OU2C8Pb1gwYwCgYIKoZIzj0EAwMDaQAwZgIxANkQhJ0dt8SOSna6XMugwcm27C39vXpGxMZpkhkGz9jtX32uMYxS7PU7bOQ8UK/J8gIxAJOXk26SNI4T6amV9Qajdke3UVMOwIWFLxjo1eq6gai1C+I5ew1b3oJ9p5MB56BxMA=="}, "tlogEntries":[{"logIndex":"2613969849", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787814093", "inclusionPromise":{"signedEntryTimestamp":"MEQCIEw/H/N+EZIwnyYubbQjbe4l1vEThipNr1UfJoGz9FDkAiBDxRJPxNCN7YSdiMTbOeTcChXz/3rXbPSz4Ggs0cq0xg=="}, "inclusionProof":{"logIndex":"2492065587", "rootHash":"HMabR7yd/EVq4WbRn7sHFJlrkYwtTQFLVw8bfn+V/3U=", "treeSize":"2492065618", "hashes":["/47HCMUHWevEtJadH20kmmpM1uwTvrnQPGl7XIxbn0g=", "S6DPOYJB7pC6yD+CDg/76n9DvCEnYP1bWwSyBOxtyBQ=", "08UjMwpJoruGrE0eG1Oa0j+eZiI3bGuJgz49Ck04JEE=", "y1zjtlmIKKMItKXBRoSrfvYznV5kx1gp1TH/JZnMvDQ=", "3QWj8I1L/XMP6bvhV6UfDbCFQtLroZAXR8elY5sR1f8=", "eXbrdCAOkgErRva+XLB62kDcJwlWngQ4NqKGch5Wz40=", "YpZEITJioo+i82eneJRh/rPljqYzpU6JkSijrmcgZZs=", "7f1DiyYwq4ZLq0MB+i2yDsDcVgCFJbIE4HvxSM1mYAo=", "+1z062t3T904zuuUCVWmzD3bmZKoTqs9bZ1VmfwnByQ=", "JZTdykpRua3w/oyS32G9vqOIoVajdstcZWLNMlPQfac=", "UbnDpZmg3ac4T8gCmplTGV9yU5j3fWOpeZi+i3apFqM=", "ygNYNZXNiY/+rjtW6HI0ngBM75qjGCTD6YstKiLiOZY=", "RvtgIftU0w4e16LAcGD7A1ErF9x2p/QLAyDLR+PaDMU=", "kW2yseLcah3rCzJJ1VS5Uhr9NCuHTJMDHsUkYy/0dV0=", "VGrFt5JvlIUQvoGEdBBxrCe3lYbu581OhgLq/Xdyat4=", "7TfDCaMWgEwSJEZd/e33b24919rF4IDJLmbqNNRNhZs=", "0HfdxDt/zGugZLIuHasrdLEW9s8OaTykSSqBEDFQXdw=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2492065618\nHMabR7yd/EVq4WbRn7sHFJlrkYwtTQFLVw8bfn+V/3U=\n\n— rekor.sigstore.dev wNI9ajBFAiEAndoa5ADD2zMzk/PVdEol6LUHALqSJ1aHkW2ZwNrNI5wCIGOxmepO1nyBeZT4ffX+mio5WU1k0ls6sfCVXThd1FB2\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJjYjNiM2YwMWMxMzA1OTQ2ZTBmMWRkMTQ3MjA5NDE5ODVlNTEyMWMxYzY1MTMzNDA2ZmIxNjdiNDdiOGRlYTIxIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUURZQkV4dHlHcm1nQXVMVEE2SVVtRmtHY0F1K2ZOd3VyTmN5SXZyYnNTc2xRSWdkRUxKcUNqQkowWGJ1UnpsQnQ0QVVNMU84NGhkTnNEWHNQa2t1MGlkZGlNPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVWRU5EUW5STFowRjNTVUpCWjBsVlpVWmhSWE5FWnpsVVRrOXJWRGhxVUdneFpVZzNXRmN4V2tjd2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1ROTlJHTjNUVlJOZWxkb1kwNU5hbGwzVDBSSk0wMUVZM2hOVkUxNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVV6VFhGc016WTNhRFZqTmpkV1FuaDJZbWhWYkhWeWFrODBWME55T0dJNGFFVndUR1FLTHpsak1VZDRXVFZVUzJoNGJTOVJXbUl6TDA1QmFXOWFTbkoxTjJ4c2FrUXZUSEUyTW1oaGNqTnpURXd3WjNveVNqWlBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZrWTJJeENrdExXV0ZaYmpZd2VGWnRkVTVOVGxKVlJrcDBUSE5SZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVpWMUpyV21wUk0wNUVSbTFaVjFwcFRsUldhVTVIVFhoTlJFWnBDazFVVlRWUFJGWnJXbTFHYkZwWFJUTlplbVJxVFcxWmVsbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWxOVkZVMVQwUldhMXB0Um13S1dsZEZNMWw2WkdwTmJWbDZXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWtLVFZSVk5VOUVWbXRhYlVac1dsZEZNMWw2WkdwTmJWbDZXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZEdhMXBIV1RBS1RucFJlRnB0Um0xWmFsVXhXV3BTYWsxVVFYaFpha1V4VDFSbk1WcEhXbWhhVjFab1RqSk5NMWw2U20xTk1rVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWsxM1RrUlZlazlVVlRWT1ZFbDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUZTVWRKTWxWQlFVRlJSQXBCUldkM1VtZEphRUZOUld3eFQzUlljbU1yY21kNlR6aEdNMFpETlc0ek5UTjFiSEEzVm1sUlpWWnJkaTk2WWt0NWVDdDZRV2xGUVhjeldXUjBRM2ROQ21ZeGRqaDNRMmx4U0VsclMzZFNPR04zUnpGb1MxWlJOazlWTWtNNFVHSXhaM2RaZDBObldVbExiMXBKZW1vd1JVRjNUVVJoVVVGM1dtZEplRUZPYTFFS2FFb3daSFE0VTA5VGJtRTJXRTExWjNkamJUSTNRek01ZGxod1IzaE5XbkJyYUd0SGVqbHFkRmd6TW5WTldYaFROMUJWTjJKUFVUaFZTeTlLT0dkSmVBcEJTazlZYXpJMlUwNUpORlEyWVcxV09WRmhhbVJyWlROVlZrMVBkMGxYUmt4NGFtOHhaWEUyWjJGcE1VTXJTVFZsZHpGaU0yOUtPWEExVFVJMU5rSjRDazFCUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgu68d7Mo1lxzqTLEcIIIYY2FmvZOIIiqmW+wOyc6iEOACFQDFjYG2FqqSK3niCsbatiEsgmH+QRgPMjAyNjA4MjcwNzAxMzNaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgyNzA3MDEzM1owLwYJKoZIhvcNAQkEMSIEINIMPF837ttfhQyR/UXYW4dIULHr9R9e+KH3gLevt0F8MIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMGP+aJ7k1spyl0u7RBSD1p+hHTP2zGhDLvU3rSRMMvfeuXeAHLq07N1/ENr6rb8JCgIwYYxKHPJ+IzAw+GokkyFSJzNhjziLIL2bw8VZnZPFzgUgyaVTUZAnjp5UXpsOBaQX"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"yzs/AcEwWUbg8d0UcglBmF5RIcHGUTNAb7FntHuN6iE="}, "signature":"MEUCIQDYBExtyGrmgAuLTA6IUmFkGcAu+fNwurNcyIvrbsSslQIgdELJqCjBJ0XbuRzlBt4AUM1O84hdNsDXsPkku0iddiM="}}

build/torch213-cxx11-xpu20260-x86_64-linux/{_megablocks_xpu_c4d0dc5.abi3.so → _megablocks_xpu_addf474.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e9e67436714e7f1ac7513960d18a05eae815c650050754fe667d1ddef6cf2210
-size 3931504

 version https://git-lfs.github.com/spec/v1
+oid sha256:40d4c6b0feb85bd2e83ca7d619f3c6a0fb79f4da6f8b44f50f70da570de5b247
+size 8288824

build/torch213-cxx11-xpu20260-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_xpu_c4d0dc5
-ops = torch.ops._megablocks_xpu_c4d0dc5
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_xpu_c4d0dc5::{op_name}"

 import torch
+from . import _megablocks_xpu_addf474
+ops = torch.ops._megablocks_xpu_addf474
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_megablocks_xpu_addf474::{op_name}"

build/torch213-cxx11-xpu20260-x86_64-linux/megablocks/__init__.py DELETED Viewed

@@ -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")))

build/torch213-cxx11-xpu20260-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,9 +1,10 @@
 {
   "name": "megablocks",
-  "id": "_megablocks_xpu_c4d0dc5",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "xpu"
   },
@@ -38,8 +39,8 @@
       "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=",
       "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=",
       "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=",
-      "_megablocks_xpu_c4d0dc5.abi3.so": "6eZ0NnFOfxrHUTlg0YoF6ugVxlAFB1T+Zn0d3vbPIhA=",
-      "_ops.py": "KZ/xutl7d4r1cvlnq+lt8PGsR1a57OQznlwnP3eVkKI=",
       "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=",
       "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=",
       "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=",
@@ -51,7 +52,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=",
@@ -96,11 +96,11 @@
   "provenance": {
     "kernel-builder": {
       "version": "0.17.0-dev0",
-      "sha": "81580bb92577f2f7228661ca2a221fc052375709",
       "dirty": false
     },
     "kernel": {
-      "sha": "c4d0dc56661badb9cc47bde647e08fe3b4bf7458",
       "dirty": false
     }
   }

 {
   "name": "megablocks",
+  "id": "_megablocks_xpu_addf474",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
+  "kernel-depends": [],
   "backend": {
     "type": "xpu"
   },
       "_layers/mpu.py": "BX61/pmqRCVeEMnzpuK8x5NpITBJyZfOawGnDegW1O8=",
       "_layers/router.py": "GKC/5//Uf43HT6zXlBTStCMTml2YXn7Jbgq9EgTgDC0=",
       "_layers/sharedexpert_registry.py": "GfwHKjbC7lBXYq9kNifp3HudTdZKSqwsFxGxFLc2krw=",
+      "_megablocks_xpu_addf474.abi3.so": "QNTGsP64W9LoPKfWGfPGoPt59Npvi0T1D3DaVw3lskc=",
+      "_ops.py": "QBFLHcpmqB9+JkTkhdpkiEE9yER4FUDrsZgGzybNoVc=",
       "_version.py": "W3l9oUrnmvfBoIy06Lv0y2ezTQhNU04ug/txHcdWJD0=",
       "backend/__init__.py": "4bunAqqjwE93nym8CigHbqVAal7UaYGdONb8lGRfrek=",
       "backend/kernels.py": "2UFNTlmSDk3S5y6UW14MJMPVVsVhmozxQ5wFDxtZHhQ=",
       "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=",
   "provenance": {
     "kernel-builder": {
       "version": "0.17.0-dev0",
+      "commit": "7d2828e6592de641b44fb8eb896719101a2ab2fb",
       "dirty": false
     },
     "kernel": {
+      "commit": "addf4741fafb55b4c101b15985dfaeea7c7c2f3a",
       "dirty": false
     }
   }

build/torch213-cxx11-xpu20260-x86_64-linux/metadata.json.sigstore CHANGED Viewed

@@ -1 +1 @@

- {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSTCCBtCgAwIBAgIUEsfCDaGiIjMTh2fWbFNDO7MtkCowCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI1MTIwNzQ3WhcNMjYwODI1MTIxNzQ3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEmPdHGofC4sgawEo7z3SJT/Ayt9ERh4bgxvOL7buMV+eCrAw2gmH3lYDRyAHRdCkSP2KQuzjQAsjJ19LZCqzw6KOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUPsC5HZTI3gn0jq+BQjq4r5S2OgswHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYzRkMGRjNTY2NjFiYWRiOWNjNDdiZGU2NDdlMDhmZTNiNGJmNzQ1ODAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGM0ZDBkYzU2NjYxYmFkYjljYzQ3YmRlNjQ3ZTA4ZmUzYjRiZjc0NTgwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzI4NDQ2ODk4OTQvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoDjRyaUAAAQDAEYwRAIgM0GM4ReR0pFWw7Wx/buzmiWJqvaZoMT/xfGlpLZRm1cCIGUMb+0iRQuknEKoQbK+5DveMo3AC6jEpy4nLUXlBadNMAoGCCqGSM49BAMDA2cAMGQCMFodrf0jX0YBZhtrHolQlOygwpaTKfIqEd1KsAEBGzEeyglfDdeCtnlJqyjC2QabTgIwMSJ9rAuw/7JMGuRwKX7SN6y10dj6bcgrWH1h1pMOpRkDviO3UcOW6D7YK0QJh6sI"}, "tlogEntries":[{"logIndex":"2583274400", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787659668", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQCqfQPWi1/v5rUMFeKZU9k687Y42OyxhJoWYlADuFTFUwIhAIRq8YTYK2Y2WqGWhf4mVDyufmMqpoud0k4fLVFBgMVq"}, "inclusionProof":{"logIndex":"2461370138", "rootHash":"sKa21DCpYU/7LZ9ORZD9C6zjkbDLdRZrxb0LEXD2oF8=", "treeSize":"2461370140", "hashes":["56CyIHBhqSmA3l4hOHXc32F/TItpTDPvieoQQ97VSao=", "y6SajW2RSxSC+qZEPkVTOda5uUFbFye//C1W3E4fOD0=", "cTavF/Ocd3ZeX1e1jnEV1te21vU3vvj+JGzwC/gStgk=", "yW+Ys91jOHBbqiKaKMiTKMYmn8iwJam2ZsB6CNDrzzs=", "BJlxWyo+0U2jxgcimrAeK1pey6FSpHWLOetp0UXRWI4=", "JsX9xGIEwFv7MFelTobmTLz4j1Jl78aSs2G7dkybw0w=", "nwobU/9yRWhsbw80xq3x/uNV1nniMWdUKAqc3UWYsJM=", "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\n2461370140\nsKa21DCpYU/7LZ9ORZD9C6zjkbDLdRZrxb0LEXD2oF8=\n\n— rekor.sigstore.dev wNI9ajBFAiBoocRiTqMBmVriZO1sD/o5GwCtMgDY+WHYB3j+0fWVLAIhAPgdJas1OQKZUc4PR5zJpIrE5p3Emu3DDxDwPi16ZReL\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI3MzUwOTg3YTRmYjA2ODdiNDYxY2EzMjJlM2NiYzVhZmY3Y2VhZGQ4ZjBkYzcyZDQ0MmU3OTQ0MzRkNzIxYjc2In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJR09BTGdFYjdMaEJsZ09CSmpZVkhwMUk2bGZIVEZ6alZjVzdJTG9mL2U0cEFpRUFueks1RGZibGF2KzVuYVdseHV1U2lxUWJ0QmN5aHE5cFhvTE8wb0t0eHVRPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRWRU5EUW5SRFowRjNTVUpCWjBsVlJYTm1RMFJoUjJsSmFrMVVhREptVjJKR1RrUlBOMDEwYTBOdmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1RGTlZFbDNUbnBSTTFkb1kwNU5hbGwzVDBSSk1VMVVTWGhPZWxFelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ0VUdSSVIyOW1RelJ6WjJGM1JXODNlak5UU2xRdlFYbDBPVVZTYURSaVozaDJUMHdLTjJKMVRWWXJaVU55UVhjeVoyMUlNMnhaUkZKNVFVaFNaRU5yVTFBeVMxRjFlbXBSUVhOcVNqRTVURnBEY1hwM05rdFBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZRYzBNMUNraGFWRWt6WjI0d2FuRXJRbEZxY1RSeU5WTXlUMmR6ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVplbEpyVFVkU2FrNVVXVEpPYWtacFdWZFNhVTlYVG1wT1JHUnBDbHBIVlRKT1JHUnNUVVJvYlZwVVRtbE9SMHB0VG5wUk1VOUVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2xhUjFVeVRrUmtiRTFFYUcwS1dsUk9hVTVIU20xT2VsRXhUMFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWw2VW10TlIxSnFUbFJaTWs1cVJtbFpWMUpwVDFkT2FrNUVaR2tLV2tkVk1rNUVaR3hOUkdodFdsUk9hVTVIU20xT2VsRXhUMFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZE5NRnBFUW1zS1dYcFZNazVxV1hoWmJVWnJXV3BzYWxsNlVUTlpiVkpzVG1wUk0xcFVRVFJhYlZWNldXcFNhVnBxWXpCT1ZHZDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtrMFRrUlJNazlFYXpSUFZGRjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUVhbEo1WVZWQlFVRlJSQXBCUlZsM1VrRkpaMDB3UjAwMFVtVlNNSEJHVjNjM1YzZ3ZZblY2YldsWFNuRjJZVnB2VFZRdmVHWkhiSEJNV2xKdE1XTkRTVWRWVFdJck1HbFNVWFZyQ201RlMyOVJZa3NyTlVSMlpVMXZNMEZETm1wRmNIazBia3hWV0d4Q1lXUk9UVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlZMEZOUjFGRFRVWnZaSEptTUdvS1dEQlpRbHBvZEhKSWIyeFJiRTk1WjNkd1lWUkxaa2x4UldReFMzTkJSVUpIZWtWbGVXZHNaa1JrWlVOMGJteEtjWGxxUXpKUllXSlVaMGwzVFZOS09RcHlRWFYzTHpkS1RVZDFVbmRMV0RkVFRqWjVNVEJrYWpaaVkyZHlWMGd4YURGd1RVOXdVbXRFZG1sUE0xVmpUMWMyUkRkWlN6QlJTbWcyYzBrS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyDADAgEAMIICvwYJKoZIhvcNAQcCoIICsDCCAqwCAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgRIyiiGiz024oUd3PbWjmpU4Y6+JvU70/lSZrsvLelY0CFF9e0UyujyckpCkQ8/3pRK36GnSgGA8yMDI2MDgyNTEyMDc0N1owAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdowggHWAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwODI1MTIwNzQ3WjAvBgkqhkiG9w0BCQQxIgQgakL2LLBSBy+oQUb+Pe25PcygEP3fQJZCV+Gtb6iAf3IwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGYwZAIwMkmYobhjvzAoBXodVkDPISEvVJD2D8VPMKig4/BXuTpd0ZUS22yF79X5tHAMOkNmAjB/7sdz/RW1HEIQ3AN1fMQwpJY7iKj5JDvrcHaXrQJZczm6BnMfClNitGhQzvaOzDY="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"c1CYek+waHtGHKMi48vFr/fOrdjw3HLUQueUQ01yG3Y="}, "signature":"MEUCIGOALgEb7LhBlgOBJjYVHp1I6lfHTFzjVcW7ILof/e4pAiEAnzK5Dfblav+5naWlxuuSiqQbtBcyhq9pXoLO0oKtxuQ="}}

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUS6NAYvKLaslno4afvbnIQguXVIwwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODI3MDcwMTMwWhcNMjYwODI3MDcxMTMwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEkkjUAxP5KcArfTKbAVpacWLLJex7TpKZBpVwxwfTsp7a5297XvEVrmqpJL9dUaB0UbqikK8EdFFibFlh0pa8ZqOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUc9wV0QOoF2EImIrrC3wic3v6GkQwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYWRkZjQ3NDFmYWZiNTViNGMxMDFiMTU5ODVkZmFlZWE3YzdjMmYzYTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGFkZGY0NzQxZmFmYjU1YjRjMTAxYjE1OTg1ZGZhZWVhN2M3YzJmM2EwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzMwNDUzOTU5NTIvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoEIGFjUAAAQDAEcwRQIhANTRF20NOmZePyHcaLZoB+Ynh9gmeXuMkalMRZCeOecnAiBxjNS1Z2+jOeWRa6XIk3XrqnK/Iq41yIUIcglNc0MpEjAKBggqhkjOPQQDAwNoADBlAjAgugwsZ4aytQwxr/7pho7GVZocMS9UWugVQfh7RRrrQVwoNVxR0XBw0DD+SihdQKcCMQDVq/kD8ugfS/ELUyJ5sRiCGAQBNNocsFu4xUsJoY0p2XigO38hOTbnf8urjL0MOas="}, "tlogEntries":[{"logIndex":"2613968932", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1787814090", "inclusionPromise":{"signedEntryTimestamp":"MEQCIGkMZUpX0r1RdUT/qCpwNzsPBREIVwYwNNL0jyh3nV4qAiBfpetY+w32PkqAfhxX3vTpFJ6TAceJVyhQfX+QPt1uoQ=="}, "inclusionProof":{"logIndex":"2492064670", "rootHash":"1h92M2dTLnvbT7MAA2OgFilyI0+Wvly08dno+84I5lE=", "treeSize":"2492064684", "hashes":["CWmlwIKdRbEbSDeboUjozDIsx0bqOB/WEfIayAqx5YU=", "6LLc+NRuoFNMAhwGV746EP4qP91BNNMTnykCearhUc0=", "H5GPRpgXkD94jzZLD2lG1o2dNa+Dp4wZ3x3quVORDxk=", "vY0nAvZzZfKX0Pm0YTeJ2Nr7dUioCo09gqKzHqFm9O0=", "Lq+5cItalyvqW82mU5fvHUlud2o6lVXhO9TYKPjr+4A=", "2o4EZw1FLXLVQFKDPjPjvYsApHjoG1wq9ybfzyk7Kkw=", "+12tCov0Gm6R26NJWDl9P7e7QCoq/wf0/Szk2i4m4ZA=", "8den7EkYRWGnAu0x2oUyxgALn1Q7Qk2YTFObIa7zPEI=", "w0py6ki7Xp/BdPISxRvuHaSG06mdZqtk6VM+YYIIbEA=", "UbnDpZmg3ac4T8gCmplTGV9yU5j3fWOpeZi+i3apFqM=", "ygNYNZXNiY/+rjtW6HI0ngBM75qjGCTD6YstKiLiOZY=", "RvtgIftU0w4e16LAcGD7A1ErF9x2p/QLAyDLR+PaDMU=", "kW2yseLcah3rCzJJ1VS5Uhr9NCuHTJMDHsUkYy/0dV0=", "VGrFt5JvlIUQvoGEdBBxrCe3lYbu581OhgLq/Xdyat4=", "7TfDCaMWgEwSJEZd/e33b24919rF4IDJLmbqNNRNhZs=", "0HfdxDt/zGugZLIuHasrdLEW9s8OaTykSSqBEDFQXdw=", "SndbMKVtcTenAkwi2JBfGzD+mhexp1qJbRIY+A1JRIU=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2492064684\n1h92M2dTLnvbT7MAA2OgFilyI0+Wvly08dno+84I5lE=\n\n— rekor.sigstore.dev wNI9ajBFAiEA8LYuNM1a+J0KgCyc6r2RNOn5HNRKCJeErkpmjrlweCgCICrwiO3l3JOK78OZTo/eGOzQwczFgUgoKhqut6MKY+a6\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJlY2UwNTcxYTZmMmYyNzQxOTUyNmQ2YjAwMTBmOGM1OGMxNTQ3NDYyNWVmZTQwZGUyNzI4M2JhMTE5ZjQ0ODZmIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUN2WW5OR2J5TndpZFNiSzd0R1BiNG1VaVNIUjR5R1FJMEtBMGZUSU1RaEJBSWdNZnhEcWVqaTMwcG1IWEJCZHk3Q2IvYlpFZXdwNkxyV1VVR0dFZ2hJQ0xvPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlV6Wk9RVmwyUzB4aGMyeHVielJoWm5aaWJrbFJaM1ZZVmtsM2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFU1ROTlJHTjNUVlJOZDFkb1kwNU5hbGwzVDBSSk0wMUVZM2hOVkUxM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZyYTJwVlFYaFFOVXRqUVhKbVZFdGlRVlp3WVdOWFRFeEtaWGczVkhCTFdrSndWbmNLZUhkbVZITndOMkUxTWprM1dIWkZWbkp0Y1hCS1REbGtWV0ZDTUZWaWNXbHJTemhGWkVaR2FXSkdiR2d3Y0dFNFduRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZqT1hkV0NqQlJUMjlHTWtWSmJVbHlja016ZDJsak0zWTJSMnRSZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVpWMUpyV21wUk0wNUVSbTFaVjFwcFRsUldhVTVIVFhoTlJFWnBDazFVVlRWUFJGWnJXbTFHYkZwWFJUTlplbVJxVFcxWmVsbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWxOVkZVMVQwUldhMXB0Um13S1dsZEZNMWw2WkdwTmJWbDZXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWxYVW10YWFsRXpUa1JHYlZsWFdtbE9WRlpwVGtkTmVFMUVSbWtLVFZSVk5VOUVWbXRhYlVac1dsZEZNMWw2WkdwTmJWbDZXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZEdhMXBIV1RBS1RucFJlRnB0Um0xWmFsVXhXV3BTYWsxVVFYaFpha1V4VDFSbk1WcEhXbWhhVjFab1RqSk5NMWw2U20xTk1rVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWsxM1RrUlZlazlVVlRWT1ZFbDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUZTVWRHYWxWQlFVRlJSQXBCUldOM1VsRkphRUZPVkZKR01qQk9UMjFhWlZCNVNHTmhURnB2UWl0WmJtZzVaMjFsV0hWTmEyRnNUVkphUTJWUFpXTnVRV2xDZUdwT1V6RmFNaXRxQ2s5bFYxSmhObGhKYXpOWWNuRnVTeTlKY1RReGVVbFZTV05uYkU1ak1FMXdSV3BCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEJaM1ZuZDNNS1dqUmhlWFJSZDNoeUx6ZHdhRzgzUjFaYWIyTk5VemxWVjNWblZsRm1hRGRTVW5KeVVWWjNiMDVXZUZJd1dFSjNNRVJFSzFOcGFHUlJTMk5EVFZGRVZncHhMMnRFT0hWblpsTXZSVXhWZVVvMWMxSnBRMGRCVVVKT1RtOWpjMFoxTkhoVmMwcHZXVEJ3TWxocFowOHpPR2hQVkdKdVpqaDFjbXBNTUUxUFlYTTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgW8NRjp7oXtREaXDIukSp7psVkQj2lN7XSpes9QY1uoICFQDK2i1Rrxz3BNQNjaKBZFvL/WPHCxgPMjAyNjA4MjcwNzAxMzBaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgyNzA3MDEzMFowLwYJKoZIhvcNAQkEMSIEIC8EDpIkUJBXewTGsG4JIfOZtYX8DWyTISXasNinPG+PMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMGDN0sYVL3TD1HrB6Wh/Hqd43ZA7KZfW2UdfAe1NSvZkbyZ/FMz43fIxpBTZA21YpQIwDNOsqeMnIf36VRzdvs0+a5qhkSa7whBgvg5SROyG67Ql1V04UjSlV31aIJhTQW/a"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"7OBXGm8vJ0GVJtawAQ+MWMFUdGJe/kDeJyg7oRn0SG8="}, "signature":"MEUCIQCvYnNGbyNwidSbK7tGPb4mUiSHR4yGQI0KA0fTIMQhBAIgMfxDqeji30pmHXBBdy7Cb/bZEewp6LrWUUGGEghICLo="}}

build/torch214-cxx11-xpu20261-x86_64-linux/__init__.py ADDED Viewed

	@@ -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",
+]

build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/__init__.py ADDED Viewed

File without changes

build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/_triton_kernels/__init__.py ADDED Viewed

File without changes

build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/_triton_kernels/gmm.py ADDED Viewed

	@@ -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")

build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/adapter.py ADDED Viewed

	@@ -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

build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/configs.py ADDED Viewed

	@@ -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}}}}

build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/gmm.py ADDED Viewed

	@@ -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

build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/utils/__init__.py ADDED Viewed

File without changes

build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/utils/_triton/__init__.py ADDED Viewed

File without changes

build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/utils/_triton/arch_info.py ADDED Viewed

	@@ -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",)

build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/utils/_triton/pid_preprocessing.py ADDED Viewed

	@@ -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

build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/utils/gmm_common.py ADDED Viewed

	@@ -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

build/torch214-cxx11-xpu20261-x86_64-linux/_grouped_gemm_triton/utils/logger.py ADDED Viewed

	@@ -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)

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/__init__.py ADDED Viewed

	@@ -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',
+]

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/activation_fn.py ADDED Viewed

	@@ -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

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/all_to_all.py ADDED Viewed

	@@ -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,
+    )

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/arguments.py ADDED Viewed

	@@ -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

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/common.py ADDED Viewed

	@@ -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

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/dmlp_registry.py ADDED Viewed

	@@ -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)

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/dmoe.py ADDED Viewed

	@@ -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)

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/gelu.py ADDED Viewed

	@@ -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,
+    )

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/glu.py ADDED Viewed

	@@ -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))

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/memory_test.py ADDED Viewed

	@@ -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)

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/mlp.py ADDED Viewed

	@@ -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)))

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/moe.py ADDED Viewed

	@@ -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

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/mpu.py ADDED Viewed

	@@ -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)

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/router.py ADDED Viewed

	@@ -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

build/torch214-cxx11-xpu20261-x86_64-linux/_layers/sharedexpert_registry.py ADDED Viewed

	@@ -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)

build/torch214-cxx11-xpu20261-x86_64-linux/_megablocks_xpu_addf474.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:6c8d62889d098114f0801bbfcf41f6212fa9d4efc9c53de340a7a2922f76aad7
+size 3939216

build/torch214-cxx11-xpu20261-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _megablocks_xpu_addf474
+ops = torch.ops._megablocks_xpu_addf474
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_megablocks_xpu_addf474::{op_name}"

build/torch214-cxx11-xpu20261-x86_64-linux/_version.py ADDED Viewed

	@@ -0,0 +1,6 @@

+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+"""The MegaBlocks Version."""
+__version__ = '0.11.0.dev0'

build/torch214-cxx11-xpu20261-x86_64-linux/backend/__init__.py ADDED Viewed

	@@ -0,0 +1,2 @@


1	+ # Copyright 2024 Databricks
2	+ # SPDX-License-Identifier: Apache-2.0

build/torch214-cxx11-xpu20261-x86_64-linux/backend/kernels.py ADDED Viewed

	@@ -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

build/torch214-cxx11-xpu20261-x86_64-linux/benchmark_util.py ADDED Viewed

	@@ -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)

build/torch214-cxx11-xpu20261-x86_64-linux/cpu_fused_moe.py ADDED Viewed

	@@ -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",
+]

build/torch214-cxx11-xpu20261-x86_64-linux/cpu_moe_cpp.py ADDED Viewed

	@@ -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"]

build/torch214-cxx11-xpu20261-x86_64-linux/grouped_gemm/__init__.py ADDED Viewed

	@@ -0,0 +1,2 @@


1	+ from . import ops
2	+ from . import backend

build/torch214-cxx11-xpu20261-x86_64-linux/grouped_gemm/backend.py ADDED Viewed

	@@ -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

build/torch214-cxx11-xpu20261-x86_64-linux/grouped_gemm/ops.py ADDED Viewed

	@@ -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)

build/torch214-cxx11-xpu20261-x86_64-linux/grouped_gemm_util.py ADDED Viewed

	@@ -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