Spaces:

kernels-community
/

kernels-benchmarks

Running

App Files Files Community

drbh HF Staff commited on Sep 24

Commit

6094336

verified ·

1 Parent(s): ed9a6af

Upload folder using huggingface_hub

Browse files

Files changed (10) hide show

megablocks/megablocks_only.html +0 -0
megablocks_yamoe/artifacts/binned_run/binned_results.json +9 -9
megablocks_yamoe/artifacts/gptoss_run/gptoss_results.json +9 -9
megablocks_yamoe/artifacts/gptoss_training_run/gptoss_training_results.json +9 -9
megablocks_yamoe/artifacts/yamoe_run/yamoe_results.json +9 -9
megablocks_yamoe/cells/__pycache__/bench_utils.cpython-311.pyc +0 -0
megablocks_yamoe/cells/__pycache__/config.cpython-311.pyc +0 -0
megablocks_yamoe/cells/megablocks_run.py +1 -1
megablocks_yamoe/megablocks_yamoe.html +73 -79
megablocks_yamoe/torch_profile.html +222 -222

megablocks/megablocks_only.html CHANGED Viewed

The diff for this file is too large to render. See raw diff

megablocks_yamoe/artifacts/binned_run/binned_results.json CHANGED Viewed

@@ -9,16 +9,16 @@
     "vary_inputs": true
   },
   "stats": {
-    "avg_ms": 35.832872119995045,
-    "min_ms": 32.58174399991276,
-    "max_ms": 40.50060700001268,
-    "std_ms": 1.694341573523051,
-    "p50_ms": 36.17695449997882,
-    "p95_ms": 38.67062735003515,
-    "p99_ms": 39.92923416996405,
     "num_iters": 50,
-    "tokens_per_s": 2790.733594145783,
-    "throughput_variance": 131.29596945634063
   },
   "output_sum": 3.97190523147583
 }

     "vary_inputs": true
   },
   "stats": {
+    "avg_ms": 36.21514258000616,
+    "min_ms": 33.172280000030696,
+    "max_ms": 38.75413800005845,
+    "std_ms": 1.401058294284512,
+    "p50_ms": 36.36444199997868,
+    "p95_ms": 38.060839599990004,
+    "p99_ms": 38.46422802999541,
     "num_iters": 50,
+    "tokens_per_s": 2761.275888368544,
+    "throughput_variance": 108.05444381816277
   },
   "output_sum": 3.97190523147583
 }

megablocks_yamoe/artifacts/gptoss_run/gptoss_results.json CHANGED Viewed

@@ -9,16 +9,16 @@
     "vary_inputs": true
   },
   "stats": {
-    "avg_ms": 46.790802699997585,
-    "min_ms": 39.03555299996242,
-    "max_ms": 50.85692799991648,
-    "std_ms": 3.250858562771192,
-    "p50_ms": 47.475618500016026,
-    "p95_ms": 50.805645549957035,
-    "p99_ms": 50.83896361993766,
     "num_iters": 50,
-    "tokens_per_s": 2137.172141310693,
-    "throughput_variance": 155.17201487457513
   },
   "output_sum": 11.53223705291748
 }

     "vary_inputs": true
   },
   "stats": {
+    "avg_ms": 45.94982444000152,
+    "min_ms": 40.76497799997014,
+    "max_ms": 52.299967999942965,
+    "std_ms": 3.623045351544196,
+    "p50_ms": 45.46925300002158,
+    "p95_ms": 51.35251775002985,
+    "p99_ms": 52.12179027996967,
     "num_iters": 50,
+    "tokens_per_s": 2176.286878540176,
+    "throughput_variance": 169.79505096491204
   },
   "output_sum": 11.53223705291748
 }

megablocks_yamoe/artifacts/gptoss_training_run/gptoss_training_results.json CHANGED Viewed

@@ -9,16 +9,16 @@
     "vary_inputs": true
   },
   "stats": {
-    "avg_ms": 45.006849599990346,
-    "min_ms": 38.83674200005771,
-    "max_ms": 49.30821800007834,
-    "std_ms": 2.893955494967115,
-    "p50_ms": 45.57549300000119,
-    "p95_ms": 48.57250854988706,
-    "p99_ms": 48.963614720073565,
     "num_iters": 50,
-    "tokens_per_s": 2221.8840218494533,
-    "throughput_variance": 147.8630259637854
   },
   "output_sum": 11.53223705291748
 }

     "vary_inputs": true
   },
   "stats": {
+    "avg_ms": 46.09780513999567,
+    "min_ms": 38.8389360000474,
+    "max_ms": 49.40391599996019,
+    "std_ms": 2.4686999934552376,
+    "p50_ms": 47.23983950003685,
+    "p95_ms": 48.725092950002136,
+    "p99_ms": 49.16830440000467,
     "num_iters": 50,
+    "tokens_per_s": 2169.300679203864,
+    "throughput_variance": 122.29861537972276
   },
   "output_sum": 11.53223705291748
 }

megablocks_yamoe/artifacts/yamoe_run/yamoe_results.json CHANGED Viewed

@@ -9,16 +9,16 @@
     "vary_inputs": true
   },
   "stats": {
-    "avg_ms": 4.2496077999999216,
-    "min_ms": 4.143714000065302,
-    "max_ms": 4.276272000083736,
-    "std_ms": 0.02026809704303406,
-    "p50_ms": 4.251974999931463,
-    "p95_ms": 4.269103000035557,
-    "p99_ms": 4.276041210073345,
     "num_iters": 50,
-    "tokens_per_s": 23531.58331458302,
-    "throughput_variance": 113.86151920477748
   },
   "output_sum": 3.97190523147583
 }

     "vary_inputs": true
   },
   "stats": {
+    "avg_ms": 4.247618279998733,
+    "min_ms": 4.12893800000802,
+    "max_ms": 4.265831999987313,
+    "std_ms": 0.020712896658640616,
+    "p50_ms": 4.251555999985612,
+    "p95_ms": 4.263803499975438,
+    "p99_ms": 4.2652827100027935,
     "num_iters": 50,
+    "tokens_per_s": 23542.605151428495,
+    "throughput_variance": 117.11531020813602
   },
   "output_sum": 3.97190523147583
 }

megablocks_yamoe/cells/__pycache__/bench_utils.cpython-311.pyc CHANGED Viewed

Binary files a/megablocks_yamoe/cells/__pycache__/bench_utils.cpython-311.pyc and b/megablocks_yamoe/cells/__pycache__/bench_utils.cpython-311.pyc differ

megablocks_yamoe/cells/__pycache__/config.cpython-311.pyc CHANGED Viewed

Binary files a/megablocks_yamoe/cells/__pycache__/config.cpython-311.pyc and b/megablocks_yamoe/cells/__pycache__/config.cpython-311.pyc differ

megablocks_yamoe/cells/megablocks_run.py CHANGED Viewed

@@ -56,7 +56,7 @@ def build_megablocks_model(device: torch.device):
     # Attach loaded expert weights to the experts container
     e = model.experts
     e.alpha = 1.702
-    e.capacity_factor = 64
     e.gate_up_proj = torch.nn.Parameter(gate_up_proj.clone().to(device))
     e.gate_up_proj_bias = torch.nn.Parameter(gate_up_proj_bias.clone().to(device))
     e.down_proj = torch.nn.Parameter(down_proj.clone().to(device))

     # Attach loaded expert weights to the experts container
     e = model.experts
     e.alpha = 1.702
+    e.capacity_factor = 32
     e.gate_up_proj = torch.nn.Parameter(gate_up_proj.clone().to(device))
     e.gate_up_proj_bias = torch.nn.Parameter(gate_up_proj_bias.clone().to(device))
     e.down_proj = torch.nn.Parameter(down_proj.clone().to(device))

megablocks_yamoe/megablocks_yamoe.html CHANGED Viewed

@@ -3722,7 +3722,7 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
 <span onclick="toggleOutput('nv')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-nv" style="cursor: default; opacity: 0.3;">▶ uv-logs</span>
 </span> |
-Cell: nv | 0.53s
  | <button class="run-btn" onclick="runCell('nv')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('nv')">Copy</button>
 <a href="cells/nv.py" target="_blank" class="raw-btn">Raw</a>
@@ -3745,7 +3745,7 @@ Cell: nv | 0.53s
 </div>
 </div>
 <div id="output-nv" class="cell-output">
-<div class="cell-stdout">Wed Sep 24 21:05:30 2025
 +-----------------------------------------------------------------------------------------+
 | NVIDIA-SMI 570.172.08             Driver Version: 570.172.08     CUDA Version: 12.8     |
 |-----------------------------------------+------------------------+----------------------+
@@ -3754,19 +3754,19 @@ Cell: nv | 0.53s
 |                                         |                        |               MIG M. |
 |=========================================+========================+======================|
 |   0  NVIDIA A10G                    On  |   00000000:00:1B.0 Off |                    0 |
-|  0%   38C    P0             46W /  300W |       0MiB /  23028MiB |      0%      Default |
 |                                         |                        |                  N/A |
 +-----------------------------------------+------------------------+----------------------+
 |   1  NVIDIA A10G                    On  |   00000000:00:1C.0 Off |                    0 |
-|  0%   37C    P0             45W /  300W |       0MiB /  23028MiB |      0%      Default |
 |                                         |                        |                  N/A |
 +-----------------------------------------+------------------------+----------------------+
 |   2  NVIDIA A10G                    On  |   00000000:00:1D.0 Off |                    0 |
-|  0%   39C    P0             47W /  300W |       0MiB /  23028MiB |      0%      Default |
 |                                         |                        |                  N/A |
 +-----------------------------------------+------------------------+----------------------+
 |   3  NVIDIA A10G                    On  |   00000000:00:1E.0 Off |                    0 |
-|  0%   38C    P0             46W /  300W |       0MiB /  23028MiB |      0%      Default |
 |                                         |                        |                  N/A |
 +-----------------------------------------+------------------------+----------------------+
@@ -3792,7 +3792,7 @@ Cell: nv | 0.53s
 <span onclick="toggleOutput('setup2')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-setup2" onclick="toggleUvLogsFromHeader('setup2')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
-Cell: setup2 | 113.64s
  | <button class="run-btn" onclick="runCell('setup2')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('setup2')">Copy</button>
 <a href="cells/setup2.py" target="_blank" class="raw-btn">Raw</a>
@@ -4050,7 +4050,7 @@ Reasoning: low
 What is Tensor Parallelism?
 &lt;|end|&gt;&lt;|start|&gt;assistant&lt;|channel|&gt;analysis&lt;|message|&gt;We need to explain what Tensor Parallelism is. It&#x27;s a concept in distributed training of large language models. It refers to splitting the weight matrices (tensors) across multiple devices. Provide details: how it works, benefits, challenges, typical frameworks, etc. Also mention difference from data parallelism, pipeline parallelism. Provide example: splitting a weight matrix across GPUs, each GPU holds a slice, compute partial results, then gather. Provide mention of communication overhead, scaling, etc. Also mention that it&#x27;s used in large models like GPT-3, Megatron-LM, DeepSpeed. Provide references. Also mention that it&#x27;s also called model parallelism. Provide explanation of how it works in practice: e.g., for a linear layer, weight matrix W of shape (out_features, in_features). In tensor parallelism, split W along out_features dimension across GPUs. Each GPU computes partial output. Then gather outputs. Provide details on how to handle bias, etc. Provide mention of &quot;tensor model parallelism&quot; vs &quot;tensor parallelism&quot; synonyms. Provide mention of &quot;tensor parallelism&quot; in Megatron-LM: splitting weight matrices across GPUs. Provide mention of &quot;tensor parallelism&quot; in DeepSpeed: &quot;ZeRO-Offload&quot; etc. Provide mention
-Generation took 31.35 seconds
 </div>
 <div class="uv-install-logs" id="uv-logs-setup2">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
@@ -4059,31 +4059,31 @@ Downloading cpython-3.13.7-linux-x86_64-gnu (download) (32.0MiB)
  Downloading cpython-3.13.7-linux-x86_64-gnu (download)
    Updating https://github.com/huggingface/transformers.git (HEAD)
     Updated https://github.com/huggingface/transformers.git (7258ea44bc0c0a425a468f66f8559d1de8c4126d)
-Downloading sympy (6.0MiB)
    Building transformers @ git+https://github.com/huggingface/transformers.git@7258ea44bc0c0a425a468f66f8559d1de8c4126d
-Downloading hf-xet (3.0MiB)
-Downloading pillow (6.3MiB)
-Downloading nvidia-nvjitlink-cu12 (37.4MiB)
 Downloading networkx (1.9MiB)
-Downloading pygments (1.2MiB)
-Downloading tokenizers (3.1MiB)
-Downloading nvidia-cusparselt-cu12 (273.9MiB)
 Downloading nvidia-cublas-cu12 (566.8MiB)
 Downloading nvidia-cufile-cu12 (1.1MiB)
-Downloading jedi (1.5MiB)
-Downloading numpy (15.9MiB)
 Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
-Downloading nvidia-cudnn-cu12 (674.0MiB)
-Downloading nvidia-curand-cu12 (60.7MiB)
 Downloading triton (148.4MiB)
-Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
-Downloading nvidia-nccl-cu12 (307.4MiB)
 Downloading nvidia-cusolver-cu12 (255.1MiB)
-Downloading nvidia-cufft-cu12 (184.2MiB)
-Downloading matplotlib (8.3MiB)
-Downloading fonttools (4.7MiB)
-Downloading nvidia-cusparse-cu12 (274.9MiB)
 Downloading kiwisolver (1.4MiB)
 Downloading torch (846.8MiB)
  Downloading nvidia-cufile-cu12
  Downloading kiwisolver
@@ -4105,38 +4105,36 @@ Downloading torch (846.8MiB)
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
- Downloading nvidia-cusparse-cu12
  Downloading nvidia-cusparselt-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
-Installed 69 packages in 550ms
 </div>
 </div>
 <div class="cell-stderr">Fetching 3 files:   0%|          | 0/3 [00:00&lt;?, ?it/s]
-Fetching 3 files:  33%|███▎      | 1/3 [00:06&lt;00:12,  6.47s/it]
-Fetching 3 files:  67%|██████▋   | 2/3 [00:07&lt;00:03,  3.37s/it]
-Fetching 3 files: 100%|██████████| 3/3 [00:07&lt;00:00,  2.56s/it]
 You are using full precision kernels, we will dequantize the model to bf16. To use the quantized model with quantization kernels, please set use_kernels=False
 Loading checkpoint shards:   0%|          | 0/3 [00:00&lt;?, ?it/s]
-Loading checkpoint shards:  33%|███▎      | 1/3 [00:02&lt;00:04,  2.34s/it]
 Loading checkpoint shards:  67%|██████▋   | 2/3 [00:04&lt;00:02,  2.25s/it]
 Loading checkpoint shards: 100%|██████████| 3/3 [00:05&lt;00:00,  1.80s/it]
 Loading checkpoint shards: 100%|██████████| 3/3 [00:05&lt;00:00,  1.93s/it]
 INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` (revision: main) for layer `MegaBlocksMoeMLP`
 Fetching 66 files:   0%|          | 0/66 [00:00&lt;?, ?it/s]
-Fetching 66 files:   2%|▏         | 1/66 [00:00&lt;00:16,  3.87it/s]
-Fetching 66 files:  14%|█▎        | 9/66 [00:00&lt;00:03, 18.15it/s]
-Fetching 66 files:  26%|██▌       | 17/66 [00:00&lt;00:02, 24.03it/s]
-Fetching 66 files:  56%|█████▌    | 37/66 [00:00&lt;00:00, 58.06it/s]
-Fetching 66 files:  71%|███████   | 47/66 [00:01&lt;00:00, 37.14it/s]
-Fetching 66 files:  85%|████████▍ | 56/66 [00:01&lt;00:00, 39.66it/s]
-Fetching 66 files:  98%|█████████▊| 65/66 [00:01&lt;00:00, 42.21it/s]
-Fetching 66 files: 100%|██████████| 66/66 [00:01&lt;00:00, 37.62it/s]
-/tmp/uvnote-run-e6cle3et/home/.cache/uv/environments-v2/setup2-adf2810b697d7b08/lib/python3.13/site-packages/kernels/layer.py:868: UserWarning:
 No kernel mapping found for layer `None`. Check if the layer name matches one of the kernels in the mapping or add the kernel you want to use to the mapping. Defaulting to original forward implementation.
   warnings.warn(
 INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` (revision: main) for layer `MegaBlocksMoeMLP`
@@ -4163,7 +4161,7 @@ INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks
 INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` (revision: main) for layer `MegaBlocksMoeMLP`
 INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` (revision: main) for layer `MegaBlocksMoeMLP`
 INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` (revision: main) for layer `MegaBlocksMoeMLP`
-/tmp/uvnote-run-e6cle3et/home/.cache/uv/environments-v2/setup2-adf2810b697d7b08/lib/python3.13/site-packages/kernels/layer.py:868: UserWarning:
 No kernel mapping found for layer `None`. Check if the layer name matches one of the kernels in the mapping or add the kernel you want to use to the mapping. Defaulting to original forward implementation.
   warnings.warn(
 INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` (revision: main) for layer `MegaBlocksMoeMLP`
@@ -4200,7 +4198,7 @@ INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks
 <span onclick="toggleOutput('setup')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-setup" onclick="toggleUvLogsFromHeader('setup')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
-Cell: setup | 108.22s
  | <button class="run-btn" onclick="runCell('setup')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('setup')">Copy</button>
 <a href="cells/setup.py" target="_blank" class="raw-btn">Raw</a>
@@ -4459,12 +4457,8 @@ Reasoning: low
 What is Tensor Parallelism?
-&lt;|end|&gt;&lt;|start|&gt;assistant&lt;|channel|&gt;analysis&lt;|message|&gt;We need to explain what Tensor Parallelism is. It&#x27;s a concept in distributed training of large language models. It refers to splitting the weight matrices (tensors) across multiple devices. Provide details: how it works, benefits, challenges, typical frameworks, etc. Also mention difference from data parallelism, pipeline parallelism. Provide example: splitting a weight matrix across GPUs, each GPU holds a slice, compute partial results, then gather. Provide mention of Megatron-LM, DeepSpeed, etc. Also mention that it&#x27;s used for very large models that don&#x27;t fit in a single GPU. Provide explanation of how it reduces memory usage and increases throughput. Also mention that it can be combined with other parallelism strategies. Provide mention of communication overhead, scaling, etc. Provide references. Provide code snippet maybe. Provide final answer.&lt;|end|&gt;&lt;|start|&gt;assistant&lt;|channel|&gt;final&lt;|message|&gt;## Tensor Parallelism – A Quick Overview
-| **Aspect** | **What it is** | **Why it matters** |
-|------------|----------------|--------------------|
-| **Definition** | Splitting a *single* weight tensor (e.g., a large matrix in a transformer layer) across multiple devices so that each device holds only a *portion* of the tensor. | Allows training of models that are
-Generation took 26.28 seconds
 </div>
 <div class="uv-install-logs" id="uv-logs-setup">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
@@ -4473,37 +4467,37 @@ Downloading cpython-3.13.7-linux-x86_64-gnu (download) (32.0MiB)
  Downloading cpython-3.13.7-linux-x86_64-gnu (download)
    Updating https://github.com/huggingface/transformers.git (HEAD)
     Updated https://github.com/huggingface/transformers.git (7258ea44bc0c0a425a468f66f8559d1de8c4126d)
-Downloading tokenizers (3.1MiB)
-Downloading nvidia-cufile-cu12 (1.1MiB)
-   Building transformers @ git+https://github.com/huggingface/transformers.git@7258ea44bc0c0a425a468f66f8559d1de8c4126d
-Downloading nvidia-curand-cu12 (60.7MiB)
-Downloading nvidia-cusparse-cu12 (274.9MiB)
-Downloading pygments (1.2MiB)
-Downloading sympy (6.0MiB)
 Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
 Downloading nvidia-nccl-cu12 (307.4MiB)
 Downloading nvidia-cusparselt-cu12 (273.9MiB)
-Downloading nvidia-nvjitlink-cu12 (37.4MiB)
 Downloading nvidia-cufft-cu12 (184.2MiB)
-Downloading networkx (1.9MiB)
 Downloading kiwisolver (1.4MiB)
-Downloading pillow (6.3MiB)
 Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
 Downloading fonttools (4.7MiB)
-Downloading hf-xet (3.0MiB)
-Downloading numpy (15.9MiB)
-Downloading matplotlib (8.3MiB)
-Downloading nvidia-cudnn-cu12 (674.0MiB)
-Downloading nvidia-cusolver-cu12 (255.1MiB)
-Downloading jedi (1.5MiB)
-Downloading nvidia-cublas-cu12 (566.8MiB)
 Downloading triton (148.4MiB)
 Downloading torch (846.8MiB)
  Downloading nvidia-cufile-cu12
  Downloading kiwisolver
  Downloading pygments
- Downloading tokenizers
  Downloading hf-xet
  Downloading networkx
  Downloading fonttools
  Downloading pillow
@@ -4519,33 +4513,33 @@ Downloading torch (846.8MiB)
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
- Downloading nvidia-cusparse-cu12
  Downloading nvidia-cusparselt-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
-Installed 69 packages in 462ms
 </div>
 </div>
 <div class="cell-stderr">Fetching 3 files:   0%|          | 0/3 [00:00&lt;?, ?it/s]
-Fetching 3 files:  33%|███▎      | 1/3 [00:07&lt;00:14,  7.36s/it]
-Fetching 3 files:  67%|██████▋   | 2/3 [00:08&lt;00:03,  3.69s/it]
-Fetching 3 files: 100%|██████████| 3/3 [00:08&lt;00:00,  2.83s/it]
 You are using full precision kernels, we will dequantize the model to bf16. To use the quantized model with quantization kernels, please set use_kernels=False
 Loading checkpoint shards:   0%|          | 0/3 [00:00&lt;?, ?it/s]
-Loading checkpoint shards:  33%|███▎      | 1/3 [00:02&lt;00:04,  2.36s/it]
-Loading checkpoint shards:  67%|██████▋   | 2/3 [00:04&lt;00:02,  2.26s/it]
 Loading checkpoint shards: 100%|██████████| 3/3 [00:05&lt;00:00,  1.80s/it]
-Loading checkpoint shards: 100%|██████████| 3/3 [00:05&lt;00:00,  1.94s/it]
 INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for layer `Yamoe`
 Fetching 6 files:   0%|          | 0/6 [00:00&lt;?, ?it/s]
-Fetching 6 files:  17%|█▋        | 1/6 [00:00&lt;00:01,  2.82it/s]
-Fetching 6 files: 100%|██████████| 6/6 [00:00&lt;00:00, 11.61it/s]
-Fetching 6 files: 100%|██████████| 6/6 [00:00&lt;00:00, 10.04it/s]
-/tmp/uvnote-run-ga2bg_po/home/.cache/uv/environments-v2/setup-1400c3ff0fc01263/lib/python3.13/site-packages/kernels/layer.py:868: UserWarning:
 No kernel mapping found for layer `None`. Check if the layer name matches one of the kernels in the mapping or add the kernel you want to use to the mapping. Defaulting to original forward implementation.
   warnings.warn(
 INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for layer `Yamoe`
@@ -4572,7 +4566,7 @@ INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for laye
 INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for layer `Yamoe`
 INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for layer `Yamoe`
 INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for layer `Yamoe`
-/tmp/uvnote-run-ga2bg_po/home/.cache/uv/environments-v2/setup-1400c3ff0fc01263/lib/python3.13/site-packages/kernels/layer.py:868: UserWarning:
 No kernel mapping found for layer `None`. Check if the layer name matches one of the kernels in the mapping or add the kernel you want to use to the mapping. Defaulting to original forward implementation.
   warnings.warn(
 INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for layer `Yamoe`

 <span onclick="toggleOutput('nv')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-nv" style="cursor: default; opacity: 0.3;">▶ uv-logs</span>
 </span> |
+Cell: nv | 0.55s
  | <button class="run-btn" onclick="runCell('nv')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('nv')">Copy</button>
 <a href="cells/nv.py" target="_blank" class="raw-btn">Raw</a>
 </div>
 </div>
 <div id="output-nv" class="cell-output">
+<div class="cell-stdout">Wed Sep 24 22:04:34 2025
 +-----------------------------------------------------------------------------------------+
 | NVIDIA-SMI 570.172.08             Driver Version: 570.172.08     CUDA Version: 12.8     |
 |-----------------------------------------+------------------------+----------------------+
 |                                         |                        |               MIG M. |
 |=========================================+========================+======================|
 |   0  NVIDIA A10G                    On  |   00000000:00:1B.0 Off |                    0 |
+|  0%   36C    P0             45W /  300W |       0MiB /  23028MiB |      0%      Default |
 |                                         |                        |                  N/A |
 +-----------------------------------------+------------------------+----------------------+
 |   1  NVIDIA A10G                    On  |   00000000:00:1C.0 Off |                    0 |
+|  0%   37C    P0             47W /  300W |       0MiB /  23028MiB |      0%      Default |
 |                                         |                        |                  N/A |
 +-----------------------------------------+------------------------+----------------------+
 |   2  NVIDIA A10G                    On  |   00000000:00:1D.0 Off |                    0 |
+|  0%   35C    P0             47W /  300W |       0MiB /  23028MiB |      0%      Default |
 |                                         |                        |                  N/A |
 +-----------------------------------------+------------------------+----------------------+
 |   3  NVIDIA A10G                    On  |   00000000:00:1E.0 Off |                    0 |
+|  0%   37C    P0             44W /  300W |       0MiB /  23028MiB |      0%      Default |
 |                                         |                        |                  N/A |
 +-----------------------------------------+------------------------+----------------------+
 <span onclick="toggleOutput('setup2')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-setup2" onclick="toggleUvLogsFromHeader('setup2')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
+Cell: setup2 | 114.03s
  | <button class="run-btn" onclick="runCell('setup2')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('setup2')">Copy</button>
 <a href="cells/setup2.py" target="_blank" class="raw-btn">Raw</a>
 What is Tensor Parallelism?
 &lt;|end|&gt;&lt;|start|&gt;assistant&lt;|channel|&gt;analysis&lt;|message|&gt;We need to explain what Tensor Parallelism is. It&#x27;s a concept in distributed training of large language models. It refers to splitting the weight matrices (tensors) across multiple devices. Provide details: how it works, benefits, challenges, typical frameworks, etc. Also mention difference from data parallelism, pipeline parallelism. Provide example: splitting a weight matrix across GPUs, each GPU holds a slice, compute partial results, then gather. Provide mention of communication overhead, scaling, etc. Also mention that it&#x27;s used in large models like GPT-3, Megatron-LM, DeepSpeed. Provide references. Also mention that it&#x27;s also called model parallelism. Provide explanation of how it works in practice: e.g., for a linear layer, weight matrix W of shape (out_features, in_features). In tensor parallelism, split W along out_features dimension across GPUs. Each GPU computes partial output. Then gather outputs. Provide details on how to handle bias, etc. Provide mention of &quot;tensor model parallelism&quot; vs &quot;tensor parallelism&quot; synonyms. Provide mention of &quot;tensor parallelism&quot; in Megatron-LM: splitting weight matrices across GPUs. Provide mention of &quot;tensor parallelism&quot; in DeepSpeed: &quot;ZeRO-Offload&quot; etc. Provide mention
+Generation took 31.36 seconds
 </div>
 <div class="uv-install-logs" id="uv-logs-setup2">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
  Downloading cpython-3.13.7-linux-x86_64-gnu (download)
    Updating https://github.com/huggingface/transformers.git (HEAD)
     Updated https://github.com/huggingface/transformers.git (7258ea44bc0c0a425a468f66f8559d1de8c4126d)
+Downloading jedi (1.5MiB)
+Downloading pygments (1.2MiB)
    Building transformers @ git+https://github.com/huggingface/transformers.git@7258ea44bc0c0a425a468f66f8559d1de8c4126d
+Downloading matplotlib (8.3MiB)
 Downloading networkx (1.9MiB)
+Downloading sympy (6.0MiB)
 Downloading nvidia-cublas-cu12 (566.8MiB)
+Downloading nvidia-nvjitlink-cu12 (37.4MiB)
+Downloading nvidia-nccl-cu12 (307.4MiB)
+Downloading hf-xet (3.0MiB)
+Downloading nvidia-cusparse-cu12 (274.9MiB)
+Downloading fonttools (4.7MiB)
 Downloading nvidia-cufile-cu12 (1.1MiB)
 Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
 Downloading triton (148.4MiB)
 Downloading nvidia-cusolver-cu12 (255.1MiB)
+Downloading tokenizers (3.1MiB)
 Downloading kiwisolver (1.4MiB)
+Downloading nvidia-curand-cu12 (60.7MiB)
+Downloading pillow (6.3MiB)
+Downloading numpy (15.9MiB)
+Downloading nvidia-cufft-cu12 (184.2MiB)
+Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
+Downloading nvidia-cudnn-cu12 (674.0MiB)
+Downloading nvidia-cusparselt-cu12 (273.9MiB)
 Downloading torch (846.8MiB)
  Downloading nvidia-cufile-cu12
  Downloading kiwisolver
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
  Downloading nvidia-cusparselt-cu12
+ Downloading nvidia-cusparse-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
+Installed 69 packages in 509ms
 </div>
 </div>
 <div class="cell-stderr">Fetching 3 files:   0%|          | 0/3 [00:00&lt;?, ?it/s]
+Fetching 3 files:  33%|███▎      | 1/3 [00:06&lt;00:12,  6.49s/it]
+Fetching 3 files:  67%|██████▋   | 2/3 [00:07&lt;00:03,  3.44s/it]
+Fetching 3 files: 100%|██████████| 3/3 [00:07&lt;00:00,  2.60s/it]
 You are using full precision kernels, we will dequantize the model to bf16. To use the quantized model with quantization kernels, please set use_kernels=False
 Loading checkpoint shards:   0%|          | 0/3 [00:00&lt;?, ?it/s]
+Loading checkpoint shards:  33%|███▎      | 1/3 [00:02&lt;00:04,  2.35s/it]
 Loading checkpoint shards:  67%|██████▋   | 2/3 [00:04&lt;00:02,  2.25s/it]
 Loading checkpoint shards: 100%|██████████| 3/3 [00:05&lt;00:00,  1.80s/it]
 Loading checkpoint shards: 100%|██████████| 3/3 [00:05&lt;00:00,  1.93s/it]
 INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` (revision: main) for layer `MegaBlocksMoeMLP`
 Fetching 66 files:   0%|          | 0/66 [00:00&lt;?, ?it/s]
+Fetching 66 files:   2%|▏         | 1/66 [00:00&lt;00:10,  6.31it/s]
+Fetching 66 files:  14%|█▎        | 9/66 [00:00&lt;00:02, 26.39it/s]
+Fetching 66 files:  26%|██▌       | 17/66 [00:01&lt;00:03, 12.42it/s]
+Fetching 66 files:  74%|███████▍  | 49/66 [00:01&lt;00:00, 45.00it/s]
+Fetching 66 files:  91%|█████████ | 60/66 [00:01&lt;00:00, 45.67it/s]
+Fetching 66 files: 100%|██████████| 66/66 [00:01&lt;00:00, 34.31it/s]
+/tmp/uvnote-run-_uergc47/home/.cache/uv/environments-v2/setup2-adf2810b697d7b08/lib/python3.13/site-packages/kernels/layer.py:868: UserWarning:
 No kernel mapping found for layer `None`. Check if the layer name matches one of the kernels in the mapping or add the kernel you want to use to the mapping. Defaulting to original forward implementation.
   warnings.warn(
 INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` (revision: main) for layer `MegaBlocksMoeMLP`
 INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` (revision: main) for layer `MegaBlocksMoeMLP`
 INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` (revision: main) for layer `MegaBlocksMoeMLP`
 INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` (revision: main) for layer `MegaBlocksMoeMLP`
+/tmp/uvnote-run-_uergc47/home/.cache/uv/environments-v2/setup2-adf2810b697d7b08/lib/python3.13/site-packages/kernels/layer.py:868: UserWarning:
 No kernel mapping found for layer `None`. Check if the layer name matches one of the kernels in the mapping or add the kernel you want to use to the mapping. Defaulting to original forward implementation.
   warnings.warn(
 INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` (revision: main) for layer `MegaBlocksMoeMLP`
 <span onclick="toggleOutput('setup')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-setup" onclick="toggleUvLogsFromHeader('setup')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
+Cell: setup | 109.23s
  | <button class="run-btn" onclick="runCell('setup')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('setup')">Copy</button>
 <a href="cells/setup.py" target="_blank" class="raw-btn">Raw</a>
 What is Tensor Parallelism?
+&lt;|end|&gt;&lt;|start|&gt;assistant&lt;|channel|&gt;analysis&lt;|message|&gt;We need to explain what Tensor Parallelism is. It&#x27;s a concept in distributed training of large language models. It refers to splitting the weight matrices (tensors) across multiple devices. Provide details: how it works, benefits, challenges, typical use cases, differences from data parallelism, pipeline parallelism, model parallelism. Provide example: splitting a fully connected layer&#x27;s weight matrix across GPUs. Provide mention of frameworks: Megatron-LM, DeepSpeed, etc. Provide explanation of how forward/backward passes are computed. Provide mention of communication overhead, scaling, etc. Provide mention of &quot;tensor parallelism&quot; as part of &quot;model parallelism&quot; but specifically splitting tensors. Provide mention of &quot;tensor parallelism&quot; in context of transformer layers: splitting attention heads, feed-forward layers. Provide mention of &quot;tensor parallelism&quot; in context of &quot;DeepSpeed ZeRO Stage 3&quot; or &quot;Megatron-LM&#x27;s tensor parallelism&quot;. Provide mention of &quot;tensor parallelism&quot; as &quot;model parallelism across the weight matrices&quot; and &quot;tensor parallelism&quot; vs &quot;pipeline parallelism&quot;. Provide mention of &quot;tensor parallelism&quot; as &quot;splitting the weight matrix across GPUs, each GPU holds a slice of the matrix, and the input is broadcasted,
+Generation took 26.26 seconds
 </div>
 <div class="uv-install-logs" id="uv-logs-setup">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
  Downloading cpython-3.13.7-linux-x86_64-gnu (download)
    Updating https://github.com/huggingface/transformers.git (HEAD)
     Updated https://github.com/huggingface/transformers.git (7258ea44bc0c0a425a468f66f8559d1de8c4126d)
 Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
+Downloading pillow (6.3MiB)
+   Building transformers @ git+https://github.com/huggingface/transformers.git@7258ea44bc0c0a425a468f66f8559d1de8c4126d
+Downloading nvidia-cublas-cu12 (566.8MiB)
 Downloading nvidia-nccl-cu12 (307.4MiB)
 Downloading nvidia-cusparselt-cu12 (273.9MiB)
+Downloading numpy (15.9MiB)
+Downloading hf-xet (3.0MiB)
+Downloading nvidia-curand-cu12 (60.7MiB)
 Downloading nvidia-cufft-cu12 (184.2MiB)
+Downloading nvidia-cusolver-cu12 (255.1MiB)
+Downloading nvidia-cudnn-cu12 (674.0MiB)
+Downloading nvidia-cufile-cu12 (1.1MiB)
+Downloading pygments (1.2MiB)
+Downloading nvidia-nvjitlink-cu12 (37.4MiB)
+Downloading jedi (1.5MiB)
+Downloading sympy (6.0MiB)
 Downloading kiwisolver (1.4MiB)
+Downloading matplotlib (8.3MiB)
+Downloading nvidia-cusparse-cu12 (274.9MiB)
+Downloading networkx (1.9MiB)
 Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
+Downloading tokenizers (3.1MiB)
 Downloading fonttools (4.7MiB)
 Downloading triton (148.4MiB)
 Downloading torch (846.8MiB)
  Downloading nvidia-cufile-cu12
  Downloading kiwisolver
  Downloading pygments
  Downloading hf-xet
+ Downloading tokenizers
  Downloading networkx
  Downloading fonttools
  Downloading pillow
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
  Downloading nvidia-cusparselt-cu12
+ Downloading nvidia-cusparse-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
+Installed 69 packages in 464ms
 </div>
 </div>
 <div class="cell-stderr">Fetching 3 files:   0%|          | 0/3 [00:00&lt;?, ?it/s]
+Fetching 3 files:  33%|███▎      | 1/3 [00:07&lt;00:14,  7.38s/it]
+Fetching 3 files:  67%|██████▋   | 2/3 [00:08&lt;00:03,  3.64s/it]
+Fetching 3 files: 100%|██████████| 3/3 [00:08&lt;00:00,  2.80s/it]
 You are using full precision kernels, we will dequantize the model to bf16. To use the quantized model with quantization kernels, please set use_kernels=False
 Loading checkpoint shards:   0%|          | 0/3 [00:00&lt;?, ?it/s]
+Loading checkpoint shards:  33%|███▎      | 1/3 [00:02&lt;00:04,  2.34s/it]
+Loading checkpoint shards:  67%|██████▋   | 2/3 [00:04&lt;00:02,  2.25s/it]
 Loading checkpoint shards: 100%|██████████| 3/3 [00:05&lt;00:00,  1.80s/it]
+Loading checkpoint shards: 100%|██████████| 3/3 [00:05&lt;00:00,  1.93s/it]
 INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for layer `Yamoe`
 Fetching 6 files:   0%|          | 0/6 [00:00&lt;?, ?it/s]
+Fetching 6 files:  17%|█▋        | 1/6 [00:00&lt;00:00,  5.44it/s]
+Fetching 6 files:  50%|█████     | 3/6 [00:00&lt;00:00,  6.96it/s]
+Fetching 6 files: 100%|██████████| 6/6 [00:00&lt;00:00, 13.54it/s]
+/tmp/uvnote-run-jc1wbhvj/home/.cache/uv/environments-v2/setup-1400c3ff0fc01263/lib/python3.13/site-packages/kernels/layer.py:868: UserWarning:
 No kernel mapping found for layer `None`. Check if the layer name matches one of the kernels in the mapping or add the kernel you want to use to the mapping. Defaulting to original forward implementation.
   warnings.warn(
 INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for layer `Yamoe`
 INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for layer `Yamoe`
 INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for layer `Yamoe`
 INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for layer `Yamoe`
+/tmp/uvnote-run-jc1wbhvj/home/.cache/uv/environments-v2/setup-1400c3ff0fc01263/lib/python3.13/site-packages/kernels/layer.py:868: UserWarning:
 No kernel mapping found for layer `None`. Check if the layer name matches one of the kernels in the mapping or add the kernel you want to use to the mapping. Defaulting to original forward implementation.
   warnings.warn(
 INFO:root:Using layer `Yamoe` from repo `drbh/yamoe` (revision: v0.3.0) for layer `Yamoe`

megablocks_yamoe/torch_profile.html CHANGED Viewed

@@ -3720,7 +3720,7 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
 <span onclick="toggleOutput('utils')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-utils" onclick="toggleUvLogsFromHeader('utils')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
-Cell: utils | deps: torch, numpy | 34.25s
  | <button class="run-btn" onclick="runCell('utils')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('utils')">Copy</button>
 <a href="cells/utils.py" target="_blank" class="raw-btn">Raw</a>
@@ -3794,24 +3794,24 @@ Cell: utils | deps: torch, numpy | 34.25s
 <div class="uv-install-logs" id="uv-logs-utils">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
-Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
 Downloading setuptools (1.1MiB)
-Downloading nvidia-cufft-cu12 (184.2MiB)
-Downloading nvidia-cusolver-cu12 (255.1MiB)
-Downloading nvidia-cufile-cu12 (1.1MiB)
-Downloading nvidia-cusparse-cu12 (274.9MiB)
 Downloading numpy (16.2MiB)
-Downloading nvidia-cudnn-cu12 (674.0MiB)
-Downloading nvidia-cublas-cu12 (566.8MiB)
 Downloading nvidia-curand-cu12 (60.7MiB)
 Downloading nvidia-nccl-cu12 (307.4MiB)
 Downloading nvidia-cusparselt-cu12 (273.9MiB)
-Downloading sympy (6.0MiB)
-Downloading nvidia-nvjitlink-cu12 (37.4MiB)
 Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
-Downloading triton (148.3MiB)
 Downloading torch (846.9MiB)
-Downloading networkx (1.9MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading networkx
@@ -3824,13 +3824,13 @@ Downloading networkx (1.9MiB)
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
- Downloading nvidia-cusparselt-cu12
  Downloading nvidia-cusparse-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
-Installed 26 packages in 446ms
 </div>
 </div>
 </div>
@@ -3843,7 +3843,7 @@ Installed 26 packages in 446ms
 <span onclick="toggleOutput('bench_utils')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-bench_utils" onclick="toggleUvLogsFromHeader('bench_utils')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
-Cell: bench_utils | deps: torch, numpy | 35.45s
  | <button class="run-btn" onclick="runCell('bench_utils')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('bench_utils')">Copy</button>
 <a href="cells/bench_utils.py" target="_blank" class="raw-btn">Raw</a>
@@ -4331,24 +4331,24 @@ Cell: bench_utils | deps: torch, numpy | 35.45s
 <div class="uv-install-logs" id="uv-logs-bench_utils">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
-Downloading numpy (16.2MiB)
-Downloading torch (846.9MiB)
-Downloading triton (148.3MiB)
-Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
-Downloading nvidia-curand-cu12 (60.7MiB)
 Downloading setuptools (1.1MiB)
-Downloading nvidia-cufft-cu12 (184.2MiB)
-Downloading nvidia-cufile-cu12 (1.1MiB)
-Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
 Downloading nvidia-nvjitlink-cu12 (37.4MiB)
-Downloading nvidia-nccl-cu12 (307.4MiB)
 Downloading sympy (6.0MiB)
-Downloading networkx (1.9MiB)
-Downloading nvidia-cudnn-cu12 (674.0MiB)
-Downloading nvidia-cusolver-cu12 (255.1MiB)
 Downloading nvidia-cusparse-cu12 (274.9MiB)
 Downloading nvidia-cusparselt-cu12 (273.9MiB)
-Downloading nvidia-cublas-cu12 (566.8MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading networkx
@@ -4367,7 +4367,7 @@ Downloading nvidia-cublas-cu12 (566.8MiB)
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
-Installed 26 packages in 445ms
 </div>
 </div>
 </div>
@@ -4381,7 +4381,7 @@ Installed 26 packages in 445ms
 <span onclick="toggleOutput('config')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-config" onclick="toggleUvLogsFromHeader('config')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
-Cell: config | deps: torch, numpy | 34.31s
  | <button class="run-btn" onclick="runCell('config')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('config')">Copy</button>
 <a href="cells/config.py" target="_blank" class="raw-btn">Raw</a>
@@ -4441,24 +4441,24 @@ Cell: config | deps: torch, numpy | 34.31s
 <div class="uv-install-logs" id="uv-logs-config">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
-Downloading sympy (6.0MiB)
-Downloading nvidia-cufile-cu12 (1.1MiB)
-Downloading nvidia-nccl-cu12 (307.4MiB)
-Downloading nvidia-cudnn-cu12 (674.0MiB)
-Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
-Downloading nvidia-cusparselt-cu12 (273.9MiB)
-Downloading torch (846.9MiB)
-Downloading networkx (1.9MiB)
 Downloading nvidia-cufft-cu12 (184.2MiB)
 Downloading setuptools (1.1MiB)
-Downloading nvidia-curand-cu12 (60.7MiB)
-Downloading nvidia-cublas-cu12 (566.8MiB)
 Downloading nvidia-nvjitlink-cu12 (37.4MiB)
 Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
-Downloading nvidia-cusparse-cu12 (274.9MiB)
-Downloading numpy (16.2MiB)
-Downloading triton (148.3MiB)
-Downloading nvidia-cusolver-cu12 (255.1MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading networkx
@@ -4471,13 +4471,13 @@ Downloading nvidia-cusolver-cu12 (255.1MiB)
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
- Downloading nvidia-cusparse-cu12
  Downloading nvidia-cusparselt-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
-Installed 26 packages in 450ms
 </div>
 </div>
 </div>
@@ -4490,7 +4490,7 @@ Installed 26 packages in 450ms
 <span onclick="toggleOutput('save_data')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-save_data" onclick="toggleUvLogsFromHeader('save_data')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
-Cell: save_data | deps: torch, numpy | 39.54s
  | <button class="run-btn" onclick="runCell('save_data')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('save_data')">Copy</button>
 <a href="cells/save_data.py" target="_blank" class="raw-btn">Raw</a>
@@ -4585,24 +4585,24 @@ Down sum: 206.729263
 <div class="uv-install-logs" id="uv-logs-save_data">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
 Downloading nvidia-cusparselt-cu12 (273.9MiB)
-Downloading nvidia-nvjitlink-cu12 (37.4MiB)
-Downloading numpy (16.2MiB)
 Downloading setuptools (1.1MiB)
-Downloading nvidia-cublas-cu12 (566.8MiB)
 Downloading nvidia-cudnn-cu12 (674.0MiB)
 Downloading nvidia-curand-cu12 (60.7MiB)
-Downloading nvidia-cusparse-cu12 (274.9MiB)
-Downloading nvidia-nccl-cu12 (307.4MiB)
-Downloading nvidia-cufft-cu12 (184.2MiB)
-Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
 Downloading nvidia-cufile-cu12 (1.1MiB)
-Downloading nvidia-cusolver-cu12 (255.1MiB)
-Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
 Downloading sympy (6.0MiB)
 Downloading torch (846.9MiB)
-Downloading networkx (1.9MiB)
-Downloading triton (148.3MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading networkx
@@ -4621,16 +4621,16 @@ Downloading triton (148.3MiB)
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
-Installed 26 packages in 446ms
 </div>
 </div>
 <div class="cell-artifacts">
 <h4>Artifacts:</h4>
-<a href="artifacts/save_data/gate_up_proj_bias.pt" class="artifact" target="_blank">gate_up_proj_bias.pt</a>
 <a href="artifacts/save_data/down_proj_bias.pt" class="artifact" target="_blank">down_proj_bias.pt</a>
 <a href="artifacts/save_data/down_proj.pt" class="artifact" target="_blank">down_proj.pt</a>
 <a href="artifacts/save_data/router_weight.pt" class="artifact" target="_blank">router_weight.pt</a>
 <a href="artifacts/save_data/router_bias.pt" class="artifact" target="_blank">router_bias.pt</a>
 <a href="artifacts/save_data/gate_up_proj.pt" class="artifact" target="_blank">gate_up_proj.pt</a>
 </div>
 </div>
@@ -4645,7 +4645,7 @@ Installed 26 packages in 446ms
 <span onclick="toggleOutput('yamoe_run')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-yamoe_run" onclick="toggleUvLogsFromHeader('yamoe_run')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
-Cell: yamoe_run | deps: torch, kernels, numpy | 39.10s
  | <button class="run-btn" onclick="runCell('yamoe_run')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('yamoe_run')">Copy</button>
 <a href="cells/yamoe_run.py" target="_blank" class="raw-btn">Raw</a>
@@ -4938,10 +4938,10 @@ Input Variation: +0.001 * iteration (deterministic)
 Warming up (10 iterations)...
 Benchmarking (50 iterations)...
-  Progress: 20% complete (avg: 4.251 ms)
-  Progress: 40% complete (avg: 4.248 ms)
-  Progress: 60% complete (avg: 4.248 ms)
-  Progress: 80% complete (avg: 4.249 ms)
 Output tensors:
   Primary: shape=(1, 100, 1152), dtype=torch.float32, device=cuda:0, range=[-0.049506, 0.054984], mean=0.000034, std=0.006508, norm=2.208791
@@ -4951,19 +4951,19 @@ Output tensors:
 Iterations: 50
 Latency Statistics:
-  Average: 4.250 ms
-  Min:     4.144 ms
-  Max:     4.276 ms
-  Std Dev: 0.020 ms
 Percentiles:
   P50 (median): 4.252 ms
-  P95:          4.269 ms
-  P99:          4.276 ms
 Throughput:
-  Tokens/sec: 23531.6
-  Std Dev:    113.9
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Saved benchmark results to yamoe_results.json
@@ -4973,25 +4973,25 @@ Output sum: 3.971905
 <div class="uv-install-logs" id="uv-logs-yamoe_run">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
-Downloading hf-xet (3.0MiB)
-Downloading nvidia-nccl-cu12 (307.4MiB)
 Downloading networkx (1.9MiB)
-Downloading nvidia-cufft-cu12 (184.2MiB)
-Downloading nvidia-cusparselt-cu12 (273.9MiB)
 Downloading setuptools (1.1MiB)
-Downloading nvidia-cudnn-cu12 (674.0MiB)
 Downloading nvidia-cufile-cu12 (1.1MiB)
 Downloading nvidia-cusolver-cu12 (255.1MiB)
-Downloading numpy (16.2MiB)
 Downloading torch (846.9MiB)
-Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
-Downloading sympy (6.0MiB)
 Downloading nvidia-curand-cu12 (60.7MiB)
-Downloading nvidia-cublas-cu12 (566.8MiB)
 Downloading nvidia-nvjitlink-cu12 (37.4MiB)
-Downloading triton (148.3MiB)
-Downloading nvidia-cusparse-cu12 (274.9MiB)
-Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
  Downloading nvidia-cufile-cu12
  Downloading hf-xet
  Downloading setuptools
@@ -5011,13 +5011,14 @@ Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
-Installed 37 packages in 454ms
 </div>
 </div>
 <div class="cell-stderr">Fetching 6 files:   0%|          | 0/6 [00:00&lt;?, ?it/s]
-Fetching 6 files:  17%|█▋        | 1/6 [00:00&lt;00:01,  3.47it/s]
-Fetching 6 files:  50%|█████     | 3/6 [00:00&lt;00:00,  4.22it/s]
-Fetching 6 files: 100%|██████████| 6/6 [00:00&lt;00:00,  8.26it/s]</div>
 <div class="cell-artifacts">
 <h4>Artifacts:</h4>
 <a href="artifacts/yamoe_run/yamoe_results.json" class="artifact" target="_blank">yamoe_results.json</a>
@@ -5034,7 +5035,7 @@ Fetching 6 files: 100%|██████████| 6/6 [00:00&lt;00:00,  8.2
 <span onclick="toggleOutput('binned_run')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-binned_run" onclick="toggleUvLogsFromHeader('binned_run')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
-Cell: binned_run | deps: torch, numpy | 39.44s
  | <button class="run-btn" onclick="runCell('binned_run')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('binned_run')">Copy</button>
 <a href="cells/binned_run.py" target="_blank" class="raw-btn">Raw</a>
@@ -5448,10 +5449,10 @@ Input Variation: +0.001 * iteration (deterministic)
 Warming up (10 iterations)...
 Benchmarking (50 iterations)...
-  Progress: 20% complete (avg: 37.889 ms)
-  Progress: 40% complete (avg: 37.238 ms)
-  Progress: 60% complete (avg: 36.997 ms)
-  Progress: 80% complete (avg: 36.387 ms)
 Output tensors:
   Primary: shape=(1, 100, 1152), dtype=torch.float32, device=cuda:0, range=[-0.049506, 0.054984], mean=0.000034, std=0.006508, norm=2.208791
@@ -5461,19 +5462,19 @@ Output tensors:
 Iterations: 50
 Latency Statistics:
-  Average: 35.833 ms
-  Min:     32.582 ms
-  Max:     40.501 ms
-  Std Dev: 1.694 ms
 Percentiles:
-  P50 (median): 36.177 ms
-  P95:          38.671 ms
-  P99:          39.929 ms
 Throughput:
-  Tokens/sec: 2790.7
-  Std Dev:    131.3
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Saved benchmark results to binned_results.json
@@ -5483,24 +5484,24 @@ Output sum: 3.971905
 <div class="uv-install-logs" id="uv-logs-binned_run">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
-Downloading nvidia-nvjitlink-cu12 (37.4MiB)
-Downloading sympy (6.0MiB)
-Downloading nvidia-cufile-cu12 (1.1MiB)
-Downloading setuptools (1.1MiB)
-Downloading numpy (16.2MiB)
 Downloading networkx (1.9MiB)
-Downloading nvidia-cudnn-cu12 (674.0MiB)
 Downloading nvidia-cusparselt-cu12 (273.9MiB)
-Downloading nvidia-cublas-cu12 (566.8MiB)
-Downloading nvidia-cusparse-cu12 (274.9MiB)
-Downloading nvidia-cusolver-cu12 (255.1MiB)
-Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
 Downloading nvidia-cufft-cu12 (184.2MiB)
-Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
-Downloading nvidia-nccl-cu12 (307.4MiB)
 Downloading nvidia-curand-cu12 (60.7MiB)
-Downloading torch (846.9MiB)
 Downloading triton (148.3MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading networkx
@@ -5513,13 +5514,13 @@ Downloading triton (148.3MiB)
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
- Downloading nvidia-cusparselt-cu12
  Downloading nvidia-cusparse-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
-Installed 26 packages in 446ms
 </div>
 </div>
 <div class="cell-artifacts">
@@ -5538,7 +5539,7 @@ Installed 26 packages in 446ms
 <span onclick="toggleOutput('gptoss_run')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-gptoss_run" onclick="toggleUvLogsFromHeader('gptoss_run')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
-Cell: gptoss_run | deps: torch, numpy | 40.46s
  | <button class="run-btn" onclick="runCell('gptoss_run')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('gptoss_run')">Copy</button>
 <a href="cells/gptoss_run.py" target="_blank" class="raw-btn">Raw</a>
@@ -5856,10 +5857,10 @@ Input Variation: +0.001 * iteration (deterministic)
 Warming up (10 iterations)...
 Benchmarking (50 iterations)...
-  Progress: 20% complete (avg: 50.504 ms)
-  Progress: 40% complete (avg: 50.045 ms)
-  Progress: 60% complete (avg: 49.107 ms)
-  Progress: 80% complete (avg: 48.012 ms)
 Output tensors:
   Primary: shape=(1, 100, 1152), dtype=torch.float32, device=cuda:0, range=[-0.064982, 0.061193], mean=0.000100, std=0.013510, norm=4.585560
@@ -5869,19 +5870,19 @@ Output tensors:
 Iterations: 50
 Latency Statistics:
-  Average: 46.791 ms
-  Min:     39.036 ms
-  Max:     50.857 ms
-  Std Dev: 3.251 ms
 Percentiles:
-  P50 (median): 47.476 ms
-  P95:          50.806 ms
-  P99:          50.839 ms
 Throughput:
-  Tokens/sec: 2137.2
-  Std Dev:    155.2
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Saved benchmark results to gptoss_results.json
@@ -5891,23 +5892,23 @@ Output sum: 11.532237
 <div class="uv-install-logs" id="uv-logs-gptoss_run">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
 Downloading setuptools (1.1MiB)
-Downloading nvidia-cusparselt-cu12 (273.9MiB)
 Downloading nvidia-curand-cu12 (60.7MiB)
-Downloading numpy (16.2MiB)
 Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
-Downloading sympy (6.0MiB)
-Downloading torch (846.9MiB)
 Downloading nvidia-nccl-cu12 (307.4MiB)
 Downloading nvidia-nvjitlink-cu12 (37.4MiB)
-Downloading nvidia-cusolver-cu12 (255.1MiB)
-Downloading nvidia-cudnn-cu12 (674.0MiB)
-Downloading nvidia-cufft-cu12 (184.2MiB)
-Downloading networkx (1.9MiB)
 Downloading nvidia-cufile-cu12 (1.1MiB)
 Downloading triton (148.3MiB)
-Downloading nvidia-cusparse-cu12 (274.9MiB)
-Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
 Downloading nvidia-cublas-cu12 (566.8MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
@@ -5921,13 +5922,13 @@ Downloading nvidia-cublas-cu12 (566.8MiB)
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
- Downloading nvidia-cusparselt-cu12
  Downloading nvidia-cusparse-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
-Installed 26 packages in 442ms
 </div>
 </div>
 <div class="cell-artifacts">
@@ -5946,7 +5947,7 @@ Installed 26 packages in 442ms
 <span onclick="toggleOutput('gptoss_training_run')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-gptoss_training_run" onclick="toggleUvLogsFromHeader('gptoss_training_run')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
-Cell: gptoss_training_run | deps: torch, numpy | 39.65s
  | <button class="run-btn" onclick="runCell('gptoss_training_run')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('gptoss_training_run')">Copy</button>
 <a href="cells/gptoss_training_run.py" target="_blank" class="raw-btn">Raw</a>
@@ -6247,10 +6248,10 @@ Input Variation: +0.001 * iteration (deterministic)
 Warming up (10 iterations)...
 Benchmarking (50 iterations)...
-  Progress: 20% complete (avg: 48.334 ms)
-  Progress: 40% complete (avg: 47.917 ms)
-  Progress: 60% complete (avg: 47.077 ms)
-  Progress: 80% complete (avg: 46.038 ms)
 Output tensors:
   Primary: shape=(1, 100, 1152), dtype=torch.float32, device=cuda:0, range=[-0.064982, 0.061193], mean=0.000100, std=0.013510, norm=4.585560
@@ -6260,19 +6261,19 @@ Output tensors:
 Iterations: 50
 Latency Statistics:
-  Average: 45.007 ms
-  Min:     38.837 ms
-  Max:     49.308 ms
-  Std Dev: 2.894 ms
 Percentiles:
-  P50 (median): 45.575 ms
-  P95:          48.573 ms
-  P99:          48.964 ms
 Throughput:
-  Tokens/sec: 2221.9
-  Std Dev:    147.9
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Saved benchmark results to gptoss_training_results.json
@@ -6282,24 +6283,24 @@ Output sum: 11.532237
 <div class="uv-install-logs" id="uv-logs-gptoss_training_run">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
-Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
-Downloading networkx (1.9MiB)
 Downloading setuptools (1.1MiB)
-Downloading numpy (16.2MiB)
 Downloading nvidia-cusparselt-cu12 (273.9MiB)
-Downloading nvidia-curand-cu12 (60.7MiB)
 Downloading nvidia-cudnn-cu12 (674.0MiB)
-Downloading sympy (6.0MiB)
 Downloading nvidia-cusolver-cu12 (255.1MiB)
-Downloading nvidia-nccl-cu12 (307.4MiB)
-Downloading nvidia-cusparse-cu12 (274.9MiB)
-Downloading nvidia-nvjitlink-cu12 (37.4MiB)
 Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
 Downloading nvidia-cufft-cu12 (184.2MiB)
-Downloading nvidia-cufile-cu12 (1.1MiB)
-Downloading nvidia-cublas-cu12 (566.8MiB)
 Downloading triton (148.3MiB)
-Downloading torch (846.9MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading networkx
@@ -6318,7 +6319,7 @@ Downloading torch (846.9MiB)
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
-Installed 26 packages in 448ms
 </div>
 </div>
 <div class="cell-artifacts">
@@ -6337,7 +6338,7 @@ Installed 26 packages in 448ms
 <span onclick="toggleOutput('megablocks_run')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-megablocks_run" onclick="toggleUvLogsFromHeader('megablocks_run')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
-Cell: megablocks_run | deps: torch, numpy, kernels | 41.38s | FAILED
  | <button class="run-btn" onclick="runCell('megablocks_run')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('megablocks_run')">Copy</button>
 <a href="cells/megablocks_run.py" target="_blank" class="raw-btn">Raw</a>
@@ -6492,7 +6493,7 @@ Cell: megablocks_run | deps: torch, numpy, kernels | 41.38s | FAILED
     <span class="c1"># Attach loaded expert weights to the experts container</span>
     <span class="n">e</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">experts</span>
     <span class="n">e</span><span class="o">.</span><span class="n">alpha</span> <span class="o">=</span> <span class="mf">1.702</span>
-    <span class="n">e</span><span class="o">.</span><span class="n">capacity_factor</span> <span class="o">=</span> <span class="mi">64</span>
     <span class="n">e</span><span class="o">.</span><span class="n">gate_up_proj</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">gate_up_proj</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">))</span>
     <span class="n">e</span><span class="o">.</span><span class="n">gate_up_proj_bias</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">gate_up_proj_bias</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">))</span>
     <span class="n">e</span><span class="o">.</span><span class="n">down_proj</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">down_proj</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">))</span>
@@ -6569,25 +6570,25 @@ Warming up (10 iterations)...
 <div class="uv-install-logs" id="uv-logs-megablocks_run">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
-Downloading nvidia-cublas-cu12 (566.8MiB)
-Downloading setuptools (1.1MiB)
 Downloading numpy (16.2MiB)
-Downloading networkx (1.9MiB)
-Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
 Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
-Downloading nvidia-cusparselt-cu12 (273.9MiB)
-Downloading nvidia-nvjitlink-cu12 (37.4MiB)
 Downloading nvidia-cufile-cu12 (1.1MiB)
-Downloading torch (846.9MiB)
-Downloading nvidia-nccl-cu12 (307.4MiB)
-Downloading nvidia-cufft-cu12 (184.2MiB)
 Downloading hf-xet (3.0MiB)
-Downloading nvidia-curand-cu12 (60.7MiB)
 Downloading nvidia-cusparse-cu12 (274.9MiB)
-Downloading nvidia-cusolver-cu12 (255.1MiB)
-Downloading triton (148.3MiB)
 Downloading nvidia-cudnn-cu12 (674.0MiB)
 Downloading sympy (6.0MiB)
  Downloading nvidia-cufile-cu12
  Downloading hf-xet
  Downloading setuptools
@@ -6601,26 +6602,25 @@ Downloading sympy (6.0MiB)
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
- Downloading nvidia-cusparselt-cu12
  Downloading nvidia-cusparse-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
-Installed 37 packages in 543ms
 </div>
 </div>
 <div class="cell-stderr">Fetching 66 files:   0%|          | 0/66 [00:00&lt;?, ?it/s]
-Fetching 66 files:   2%|▏         | 1/66 [00:00&lt;00:27,  2.39it/s]
-Fetching 66 files:   6%|▌         | 4/66 [00:00&lt;00:07,  8.04it/s]
-Fetching 66 files:  17%|█▋        | 11/66 [00:00&lt;00:02, 21.45it/s]
-Fetching 66 files:  26%|██▌       | 17/66 [00:01&lt;00:02, 17.15it/s]
-Fetching 66 files:  48%|████▊     | 32/66 [00:01&lt;00:01, 30.72it/s]
-Fetching 66 files:  62%|██████▏   | 41/66 [00:01&lt;00:01, 23.83it/s]
-Fetching 66 files:  71%|███████   | 47/66 [00:02&lt;00:00, 25.88it/s]
-Fetching 66 files: 100%|██████████| 66/66 [00:02&lt;00:00, 45.13it/s]
-Fetching 66 files: 100%|██████████| 66/66 [00:02&lt;00:00, 29.34it/s]
-/tmp/tmpq5pei8xr/cuda_utils.c:5:10: fatal error: Python.h: No such file or directory
     5 | #include &lt;Python.h&gt;
       |          ^~~~~~~~~~
 compilation terminated.
@@ -6637,87 +6637,87 @@ Traceback (most recent call last):
   File &quot;/repo/moe_benchmarks/megablocks_yamoe/.uvnote/cells/bench_utils.py&quot;, line 177, in &lt;lambda&gt;
     call = lambda x: fn(x, *args[1:], **kwargs)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/torch/nn/modules/module.py&quot;, line 1773, in _wrapped_call_impl
     return self._call_impl(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/torch/nn/modules/module.py&quot;, line 1784, in _call_impl
     return forward_call(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File &quot;/repo/moe_benchmarks/megablocks_yamoe/.uvnote/cells/megablocks_run.py&quot;, line 81, in forward
     output, dummy_routing_weights = self.model(hidden_states)
                                     ^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/torch/nn/modules/module.py&quot;, line 1773, in _wrapped_call_impl
     return self._call_impl(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/torch/nn/modules/module.py&quot;, line 1784, in _call_impl
     return forward_call(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/layers.py&quot;, line 896, in forward
     output, expert_weights_out, *_ = moe_forward(
                                      ^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/layers.py&quot;, line 730, in moe_forward
     x, tokens_per_expert = forward_fn(**forward_args)
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/layers.py&quot;, line 457, in forward_once
     x = permute_and_compute(
         ^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/layers.py&quot;, line 401, in permute_and_compute
     x = ops.binned_gather(x, indices, bins, expert_capacity, top_k)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/torch/autograd/function.py&quot;, line 576, in apply
     return super().apply(*args, **kwargs)  # type: ignore[misc]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/ops/stk_autocast.py&quot;, line 30, in decorate_fwd
     return fwd(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/ops/binned_gather.py&quot;, line 26, in forward
     return kernels.binned_gather(x, indices, None, bins, bin_size, top_k)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/backend/kernels.py&quot;, line 419, in binned_gather
     _binned_copy[(num_experts, expert_capacity)](
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/jit.py&quot;, line 390, in &lt;lambda&gt;
     return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/autotuner.py&quot;, line 239, in run
     benchmark()
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/autotuner.py&quot;, line 228, in benchmark
     timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/autotuner.py&quot;, line 228, in &lt;dictcomp&gt;
     timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/autotuner.py&quot;, line 160, in _bench
     return self.do_bench(kernel_call, quantiles=(0.5, 0.2, 0.8))
            ^^^^^^^^^^^^^
   File &quot;/usr/lib/python3.11/functools.py&quot;, line 1001, in __get__
     val = self.func(instance)
           ^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/autotuner.py&quot;, line 121, in do_bench
     return driver.active.get_benchmarker()
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/driver.py&quot;, line 30, in __getattr__
     return getattr(self._initialize_obj(), name)
                    ^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/driver.py&quot;, line 26, in _initialize_obj
     self._obj = self._init_fn()
                 ^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/driver.py&quot;, line 12, in _create_driver
     return active_drivers[0]()
            ^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/backends/nvidia/driver.py&quot;, line 715, in __init__
     self.utils = CudaUtils()  # TODO: make static
                  ^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/backends/nvidia/driver.py&quot;, line 62, in __init__
     mod = compile_module_from_src(
           ^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/build.py&quot;, line 88, in compile_module_from_src
     so = _build(name, src_path, tmpdir, library_dirs or [], include_dirs or [], libraries or [])
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-  File &quot;/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/build.py&quot;, line 51, in _build
     subprocess.check_call(cc_cmd, stdout=subprocess.DEVNULL)
   File &quot;/usr/lib/python3.11/subprocess.py&quot;, line 413, in check_call
     raise CalledProcessError(retcode, cmd)
-subprocess.CalledProcessError: Command &#x27;[&#x27;/usr/bin/gcc&#x27;, &#x27;/tmp/tmpq5pei8xr/cuda_utils.c&#x27;, &#x27;-O3&#x27;, &#x27;-shared&#x27;, &#x27;-fPIC&#x27;, &#x27;-Wno-psabi&#x27;, &#x27;-o&#x27;, &#x27;/tmp/tmpq5pei8xr/cuda_utils.cpython-311-x86_64-linux-gnu.so&#x27;, &#x27;-lcuda&#x27;, &#x27;-L/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/backends/nvidia/lib&#x27;, &#x27;-L/usr/lib/x86_64-linux-gnu&#x27;, &#x27;-I/tmp/uvnote-run-ab5uowvg/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/backends/nvidia/include&#x27;, &#x27;-I/tmp/tmpq5pei8xr&#x27;, &#x27;-I/usr/include/python3.11&#x27;]&#x27; returned non-zero exit status 1.</div>
 </div>
 </div>

 <span onclick="toggleOutput('utils')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-utils" onclick="toggleUvLogsFromHeader('utils')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
+Cell: utils | deps: torch, numpy | 35.29s
  | <button class="run-btn" onclick="runCell('utils')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('utils')">Copy</button>
 <a href="cells/utils.py" target="_blank" class="raw-btn">Raw</a>
 <div class="uv-install-logs" id="uv-logs-utils">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
+Downloading networkx (1.9MiB)
 Downloading setuptools (1.1MiB)
 Downloading numpy (16.2MiB)
+Downloading sympy (6.0MiB)
 Downloading nvidia-curand-cu12 (60.7MiB)
+Downloading nvidia-nvjitlink-cu12 (37.4MiB)
+Downloading nvidia-cusparse-cu12 (274.9MiB)
+Downloading nvidia-cusolver-cu12 (255.1MiB)
+Downloading nvidia-cudnn-cu12 (674.0MiB)
 Downloading nvidia-nccl-cu12 (307.4MiB)
+Downloading nvidia-cufile-cu12 (1.1MiB)
 Downloading nvidia-cusparselt-cu12 (273.9MiB)
+Downloading nvidia-cublas-cu12 (566.8MiB)
+Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
+Downloading nvidia-cufft-cu12 (184.2MiB)
 Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
 Downloading torch (846.9MiB)
+Downloading triton (148.3MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading networkx
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
  Downloading nvidia-cusparse-cu12
+ Downloading nvidia-cusparselt-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
+Installed 26 packages in 455ms
 </div>
 </div>
 </div>
 <span onclick="toggleOutput('bench_utils')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-bench_utils" onclick="toggleUvLogsFromHeader('bench_utils')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
+Cell: bench_utils | deps: torch, numpy | 34.44s
  | <button class="run-btn" onclick="runCell('bench_utils')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('bench_utils')">Copy</button>
 <a href="cells/bench_utils.py" target="_blank" class="raw-btn">Raw</a>
 <div class="uv-install-logs" id="uv-logs-bench_utils">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
 Downloading setuptools (1.1MiB)
+Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
+Downloading nvidia-cublas-cu12 (566.8MiB)
+Downloading nvidia-cudnn-cu12 (674.0MiB)
 Downloading nvidia-nvjitlink-cu12 (37.4MiB)
+Downloading nvidia-cufile-cu12 (1.1MiB)
 Downloading sympy (6.0MiB)
 Downloading nvidia-cusparse-cu12 (274.9MiB)
 Downloading nvidia-cusparselt-cu12 (273.9MiB)
+Downloading triton (148.3MiB)
+Downloading nvidia-curand-cu12 (60.7MiB)
+Downloading torch (846.9MiB)
+Downloading networkx (1.9MiB)
+Downloading nvidia-cusolver-cu12 (255.1MiB)
+Downloading nvidia-cufft-cu12 (184.2MiB)
+Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
+Downloading numpy (16.2MiB)
+Downloading nvidia-nccl-cu12 (307.4MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading networkx
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
+Installed 26 packages in 447ms
 </div>
 </div>
 </div>
 <span onclick="toggleOutput('config')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-config" onclick="toggleUvLogsFromHeader('config')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
+Cell: config | deps: torch, numpy | 34.69s
  | <button class="run-btn" onclick="runCell('config')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('config')">Copy</button>
 <a href="cells/config.py" target="_blank" class="raw-btn">Raw</a>
 <div class="uv-install-logs" id="uv-logs-config">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
+Downloading numpy (16.2MiB)
 Downloading nvidia-cufft-cu12 (184.2MiB)
+Downloading nvidia-cusolver-cu12 (255.1MiB)
+Downloading torch (846.9MiB)
+Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
 Downloading setuptools (1.1MiB)
+Downloading triton (148.3MiB)
+Downloading nvidia-cusparse-cu12 (274.9MiB)
+Downloading networkx (1.9MiB)
+Downloading nvidia-cusparselt-cu12 (273.9MiB)
+Downloading nvidia-nccl-cu12 (307.4MiB)
 Downloading nvidia-nvjitlink-cu12 (37.4MiB)
+Downloading sympy (6.0MiB)
+Downloading nvidia-cudnn-cu12 (674.0MiB)
+Downloading nvidia-curand-cu12 (60.7MiB)
 Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
+Downloading nvidia-cufile-cu12 (1.1MiB)
+Downloading nvidia-cublas-cu12 (566.8MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading networkx
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
  Downloading nvidia-cusparselt-cu12
+ Downloading nvidia-cusparse-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
+Installed 26 packages in 526ms
 </div>
 </div>
 </div>
 <span onclick="toggleOutput('save_data')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-save_data" onclick="toggleUvLogsFromHeader('save_data')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
+Cell: save_data | deps: torch, numpy | 40.40s
  | <button class="run-btn" onclick="runCell('save_data')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('save_data')">Copy</button>
 <a href="cells/save_data.py" target="_blank" class="raw-btn">Raw</a>
 <div class="uv-install-logs" id="uv-logs-save_data">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
+Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
+Downloading nvidia-nccl-cu12 (307.4MiB)
+Downloading nvidia-cusparse-cu12 (274.9MiB)
 Downloading nvidia-cusparselt-cu12 (273.9MiB)
 Downloading setuptools (1.1MiB)
 Downloading nvidia-cudnn-cu12 (674.0MiB)
+Downloading numpy (16.2MiB)
+Downloading triton (148.3MiB)
+Downloading networkx (1.9MiB)
+Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
 Downloading nvidia-curand-cu12 (60.7MiB)
 Downloading nvidia-cufile-cu12 (1.1MiB)
+Downloading nvidia-cufft-cu12 (184.2MiB)
 Downloading sympy (6.0MiB)
 Downloading torch (846.9MiB)
+Downloading nvidia-cublas-cu12 (566.8MiB)
+Downloading nvidia-cusolver-cu12 (255.1MiB)
+Downloading nvidia-nvjitlink-cu12 (37.4MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading networkx
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
+Installed 26 packages in 563ms
 </div>
 </div>
 <div class="cell-artifacts">
 <h4>Artifacts:</h4>
 <a href="artifacts/save_data/down_proj_bias.pt" class="artifact" target="_blank">down_proj_bias.pt</a>
 <a href="artifacts/save_data/down_proj.pt" class="artifact" target="_blank">down_proj.pt</a>
 <a href="artifacts/save_data/router_weight.pt" class="artifact" target="_blank">router_weight.pt</a>
 <a href="artifacts/save_data/router_bias.pt" class="artifact" target="_blank">router_bias.pt</a>
+<a href="artifacts/save_data/gate_up_proj_bias.pt" class="artifact" target="_blank">gate_up_proj_bias.pt</a>
 <a href="artifacts/save_data/gate_up_proj.pt" class="artifact" target="_blank">gate_up_proj.pt</a>
 </div>
 </div>
 <span onclick="toggleOutput('yamoe_run')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-yamoe_run" onclick="toggleUvLogsFromHeader('yamoe_run')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
+Cell: yamoe_run | deps: torch, kernels, numpy | 38.77s
  | <button class="run-btn" onclick="runCell('yamoe_run')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('yamoe_run')">Copy</button>
 <a href="cells/yamoe_run.py" target="_blank" class="raw-btn">Raw</a>
 Warming up (10 iterations)...
 Benchmarking (50 iterations)...
+  Progress: 20% complete (avg: 4.248 ms)
+  Progress: 40% complete (avg: 4.246 ms)
+  Progress: 60% complete (avg: 4.247 ms)
+  Progress: 80% complete (avg: 4.247 ms)
 Output tensors:
   Primary: shape=(1, 100, 1152), dtype=torch.float32, device=cuda:0, range=[-0.049506, 0.054984], mean=0.000034, std=0.006508, norm=2.208791
 Iterations: 50
 Latency Statistics:
+  Average: 4.248 ms
+  Min:     4.129 ms
+  Max:     4.266 ms
+  Std Dev: 0.021 ms
 Percentiles:
   P50 (median): 4.252 ms
+  P95:          4.264 ms
+  P99:          4.265 ms
 Throughput:
+  Tokens/sec: 23542.6
+  Std Dev:    117.1
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Saved benchmark results to yamoe_results.json
 <div class="uv-install-logs" id="uv-logs-yamoe_run">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
 Downloading networkx (1.9MiB)
+Downloading sympy (6.0MiB)
 Downloading setuptools (1.1MiB)
+Downloading nvidia-nccl-cu12 (307.4MiB)
+Downloading hf-xet (3.0MiB)
 Downloading nvidia-cufile-cu12 (1.1MiB)
+Downloading triton (148.3MiB)
+Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
+Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
 Downloading nvidia-cusolver-cu12 (255.1MiB)
+Downloading nvidia-cusparse-cu12 (274.9MiB)
+Downloading nvidia-cudnn-cu12 (674.0MiB)
+Downloading nvidia-cufft-cu12 (184.2MiB)
 Downloading torch (846.9MiB)
 Downloading nvidia-curand-cu12 (60.7MiB)
 Downloading nvidia-nvjitlink-cu12 (37.4MiB)
+Downloading nvidia-cusparselt-cu12 (273.9MiB)
+Downloading nvidia-cublas-cu12 (566.8MiB)
+Downloading numpy (16.2MiB)
  Downloading nvidia-cufile-cu12
  Downloading hf-xet
  Downloading setuptools
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
+Installed 37 packages in 449ms
 </div>
 </div>
 <div class="cell-stderr">Fetching 6 files:   0%|          | 0/6 [00:00&lt;?, ?it/s]
+Fetching 6 files:  17%|█▋        | 1/6 [00:00&lt;00:00,  5.90it/s]
+Fetching 6 files:  33%|███▎      | 2/6 [00:00&lt;00:00,  7.70it/s]
+Fetching 6 files:  50%|█████     | 3/6 [00:00&lt;00:00,  4.70it/s]
+Fetching 6 files: 100%|██████████| 6/6 [00:00&lt;00:00, 10.28it/s]</div>
 <div class="cell-artifacts">
 <h4>Artifacts:</h4>
 <a href="artifacts/yamoe_run/yamoe_results.json" class="artifact" target="_blank">yamoe_results.json</a>
 <span onclick="toggleOutput('binned_run')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-binned_run" onclick="toggleUvLogsFromHeader('binned_run')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
+Cell: binned_run | deps: torch, numpy | 38.76s
  | <button class="run-btn" onclick="runCell('binned_run')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('binned_run')">Copy</button>
 <a href="cells/binned_run.py" target="_blank" class="raw-btn">Raw</a>
 Warming up (10 iterations)...
 Benchmarking (50 iterations)...
+  Progress: 20% complete (avg: 37.794 ms)
+  Progress: 40% complete (avg: 37.656 ms)
+  Progress: 60% complete (avg: 37.188 ms)
+  Progress: 80% complete (avg: 36.704 ms)
 Output tensors:
   Primary: shape=(1, 100, 1152), dtype=torch.float32, device=cuda:0, range=[-0.049506, 0.054984], mean=0.000034, std=0.006508, norm=2.208791
 Iterations: 50
 Latency Statistics:
+  Average: 36.215 ms
+  Min:     33.172 ms
+  Max:     38.754 ms
+  Std Dev: 1.401 ms
 Percentiles:
+  P50 (median): 36.364 ms
+  P95:          38.061 ms
+  P99:          38.464 ms
 Throughput:
+  Tokens/sec: 2761.3
+  Std Dev:    108.1
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Saved benchmark results to binned_results.json
 <div class="uv-install-logs" id="uv-logs-binned_run">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
 Downloading networkx (1.9MiB)
+Downloading numpy (16.2MiB)
+Downloading setuptools (1.1MiB)
 Downloading nvidia-cusparselt-cu12 (273.9MiB)
 Downloading nvidia-cufft-cu12 (184.2MiB)
+Downloading nvidia-cudnn-cu12 (674.0MiB)
+Downloading nvidia-cufile-cu12 (1.1MiB)
 Downloading nvidia-curand-cu12 (60.7MiB)
+Downloading nvidia-nccl-cu12 (307.4MiB)
+Downloading nvidia-nvjitlink-cu12 (37.4MiB)
+Downloading nvidia-cublas-cu12 (566.8MiB)
 Downloading triton (148.3MiB)
+Downloading nvidia-cusparse-cu12 (274.9MiB)
+Downloading torch (846.9MiB)
+Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
+Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
+Downloading sympy (6.0MiB)
+Downloading nvidia-cusolver-cu12 (255.1MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading networkx
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
  Downloading nvidia-cusparse-cu12
+ Downloading nvidia-cusparselt-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
+Installed 26 packages in 455ms
 </div>
 </div>
 <div class="cell-artifacts">
 <span onclick="toggleOutput('gptoss_run')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-gptoss_run" onclick="toggleUvLogsFromHeader('gptoss_run')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
+Cell: gptoss_run | deps: torch, numpy | 39.76s
  | <button class="run-btn" onclick="runCell('gptoss_run')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('gptoss_run')">Copy</button>
 <a href="cells/gptoss_run.py" target="_blank" class="raw-btn">Raw</a>
 Warming up (10 iterations)...
 Benchmarking (50 iterations)...
+  Progress: 20% complete (avg: 51.012 ms)
+  Progress: 40% complete (avg: 49.954 ms)
+  Progress: 60% complete (avg: 48.390 ms)
+  Progress: 80% complete (avg: 46.993 ms)
 Output tensors:
   Primary: shape=(1, 100, 1152), dtype=torch.float32, device=cuda:0, range=[-0.064982, 0.061193], mean=0.000100, std=0.013510, norm=4.585560
 Iterations: 50
 Latency Statistics:
+  Average: 45.950 ms
+  Min:     40.765 ms
+  Max:     52.300 ms
+  Std Dev: 3.623 ms
 Percentiles:
+  P50 (median): 45.469 ms
+  P95:          51.353 ms
+  P99:          52.122 ms
 Throughput:
+  Tokens/sec: 2176.3
+  Std Dev:    169.8
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Saved benchmark results to gptoss_results.json
 <div class="uv-install-logs" id="uv-logs-gptoss_run">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
+Downloading numpy (16.2MiB)
+Downloading networkx (1.9MiB)
+Downloading nvidia-cusparse-cu12 (274.9MiB)
 Downloading setuptools (1.1MiB)
+Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
 Downloading nvidia-curand-cu12 (60.7MiB)
 Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
 Downloading nvidia-nccl-cu12 (307.4MiB)
+Downloading nvidia-cusparselt-cu12 (273.9MiB)
 Downloading nvidia-nvjitlink-cu12 (37.4MiB)
 Downloading nvidia-cufile-cu12 (1.1MiB)
+Downloading nvidia-cufft-cu12 (184.2MiB)
 Downloading triton (148.3MiB)
+Downloading nvidia-cusolver-cu12 (255.1MiB)
+Downloading nvidia-cudnn-cu12 (674.0MiB)
+Downloading torch (846.9MiB)
+Downloading sympy (6.0MiB)
 Downloading nvidia-cublas-cu12 (566.8MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
  Downloading nvidia-cusparse-cu12
+ Downloading nvidia-cusparselt-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
+Installed 26 packages in 524ms
 </div>
 </div>
 <div class="cell-artifacts">
 <span onclick="toggleOutput('gptoss_training_run')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-gptoss_training_run" onclick="toggleUvLogsFromHeader('gptoss_training_run')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
+Cell: gptoss_training_run | deps: torch, numpy | 40.42s
  | <button class="run-btn" onclick="runCell('gptoss_training_run')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('gptoss_training_run')">Copy</button>
 <a href="cells/gptoss_training_run.py" target="_blank" class="raw-btn">Raw</a>
 Warming up (10 iterations)...
 Benchmarking (50 iterations)...
+  Progress: 20% complete (avg: 48.387 ms)
+  Progress: 40% complete (avg: 48.249 ms)
+  Progress: 60% complete (avg: 47.887 ms)
+  Progress: 80% complete (avg: 47.011 ms)
 Output tensors:
   Primary: shape=(1, 100, 1152), dtype=torch.float32, device=cuda:0, range=[-0.064982, 0.061193], mean=0.000100, std=0.013510, norm=4.585560
 Iterations: 50
 Latency Statistics:
+  Average: 46.098 ms
+  Min:     38.839 ms
+  Max:     49.404 ms
+  Std Dev: 2.469 ms
 Percentiles:
+  P50 (median): 47.240 ms
+  P95:          48.725 ms
+  P99:          49.168 ms
 Throughput:
+  Tokens/sec: 2169.3
+  Std Dev:    122.3
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Saved benchmark results to gptoss_training_results.json
 <div class="uv-install-logs" id="uv-logs-gptoss_training_run">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
+Downloading nvidia-cufile-cu12 (1.1MiB)
 Downloading setuptools (1.1MiB)
+Downloading nvidia-nvjitlink-cu12 (37.4MiB)
+Downloading sympy (6.0MiB)
 Downloading nvidia-cusparselt-cu12 (273.9MiB)
+Downloading nvidia-cublas-cu12 (566.8MiB)
+Downloading nvidia-nccl-cu12 (307.4MiB)
+Downloading torch (846.9MiB)
 Downloading nvidia-cudnn-cu12 (674.0MiB)
+Downloading networkx (1.9MiB)
+Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
+Downloading nvidia-curand-cu12 (60.7MiB)
 Downloading nvidia-cusolver-cu12 (255.1MiB)
 Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
+Downloading nvidia-cusparse-cu12 (274.9MiB)
+Downloading numpy (16.2MiB)
 Downloading nvidia-cufft-cu12 (184.2MiB)
 Downloading triton (148.3MiB)
  Downloading nvidia-cufile-cu12
  Downloading setuptools
  Downloading networkx
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
+Installed 26 packages in 451ms
 </div>
 </div>
 <div class="cell-artifacts">
 <span onclick="toggleOutput('megablocks_run')" style="cursor: pointer;">▼ output</span>
  <span id="uv-indicator-megablocks_run" onclick="toggleUvLogsFromHeader('megablocks_run')" style="cursor: pointer;">▶ uv-logs</span>
 </span> |
+Cell: megablocks_run | deps: torch, numpy, kernels | 40.19s | FAILED
  | <button class="run-btn" onclick="runCell('megablocks_run')">▶ run</button>
 <button class="copy-btn" onclick="copyCell('megablocks_run')">Copy</button>
 <a href="cells/megablocks_run.py" target="_blank" class="raw-btn">Raw</a>
     <span class="c1"># Attach loaded expert weights to the experts container</span>
     <span class="n">e</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">experts</span>
     <span class="n">e</span><span class="o">.</span><span class="n">alpha</span> <span class="o">=</span> <span class="mf">1.702</span>
+    <span class="n">e</span><span class="o">.</span><span class="n">capacity_factor</span> <span class="o">=</span> <span class="mi">32</span>
     <span class="n">e</span><span class="o">.</span><span class="n">gate_up_proj</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">gate_up_proj</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">))</span>
     <span class="n">e</span><span class="o">.</span><span class="n">gate_up_proj_bias</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">gate_up_proj_bias</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">))</span>
     <span class="n">e</span><span class="o">.</span><span class="n">down_proj</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">down_proj</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">))</span>
 <div class="uv-install-logs" id="uv-logs-megablocks_run">
 <div class="uv-logs-header" onclick="toggleUvLogs(this)">▶ UV Install Logs</div>
 <div class="uv-logs-content" style="display: none;">
 Downloading numpy (16.2MiB)
 Downloading nvidia-cuda-cupti-cu12 (9.8MiB)
 Downloading nvidia-cufile-cu12 (1.1MiB)
 Downloading hf-xet (3.0MiB)
+Downloading networkx (1.9MiB)
+Downloading torch (846.9MiB)
 Downloading nvidia-cusparse-cu12 (274.9MiB)
+Downloading nvidia-cusparselt-cu12 (273.9MiB)
 Downloading nvidia-cudnn-cu12 (674.0MiB)
+Downloading triton (148.3MiB)
+Downloading nvidia-cuda-nvrtc-cu12 (84.0MiB)
+Downloading nvidia-curand-cu12 (60.7MiB)
+Downloading nvidia-nvjitlink-cu12 (37.4MiB)
 Downloading sympy (6.0MiB)
+Downloading nvidia-cusolver-cu12 (255.1MiB)
+Downloading nvidia-cublas-cu12 (566.8MiB)
+Downloading nvidia-cufft-cu12 (184.2MiB)
+Downloading nvidia-nccl-cu12 (307.4MiB)
+Downloading setuptools (1.1MiB)
  Downloading nvidia-cufile-cu12
  Downloading hf-xet
  Downloading setuptools
  Downloading triton
  Downloading nvidia-cufft-cu12
  Downloading nvidia-cusolver-cu12
  Downloading nvidia-cusparse-cu12
+ Downloading nvidia-cusparselt-cu12
  Downloading nvidia-nccl-cu12
  Downloading nvidia-cublas-cu12
  Downloading nvidia-cudnn-cu12
  Downloading torch
+Installed 37 packages in 449ms
 </div>
 </div>
 <div class="cell-stderr">Fetching 66 files:   0%|          | 0/66 [00:00&lt;?, ?it/s]
+Fetching 66 files:   2%|▏         | 1/66 [00:00&lt;00:23,  2.74it/s]
+Fetching 66 files:  14%|█▎        | 9/66 [00:00&lt;00:03, 17.38it/s]
+Fetching 66 files:  26%|██▌       | 17/66 [00:01&lt;00:02, 17.85it/s]
+Fetching 66 files:  55%|█████▍    | 36/66 [00:01&lt;00:00, 42.23it/s]
+Fetching 66 files:  65%|██████▌   | 43/66 [00:01&lt;00:00, 38.03it/s]
+Fetching 66 files:  74%|███████▍  | 49/66 [00:01&lt;00:00, 30.77it/s]
+Fetching 66 files:  97%|█████████▋| 64/66 [00:01&lt;00:00, 48.18it/s]
+Fetching 66 files: 100%|██████████| 66/66 [00:01&lt;00:00, 34.40it/s]
+/tmp/tmptrubhjfl/cuda_utils.c:5:10: fatal error: Python.h: No such file or directory
     5 | #include &lt;Python.h&gt;
       |          ^~~~~~~~~~
 compilation terminated.
   File &quot;/repo/moe_benchmarks/megablocks_yamoe/.uvnote/cells/bench_utils.py&quot;, line 177, in &lt;lambda&gt;
     call = lambda x: fn(x, *args[1:], **kwargs)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/torch/nn/modules/module.py&quot;, line 1773, in _wrapped_call_impl
     return self._call_impl(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/torch/nn/modules/module.py&quot;, line 1784, in _call_impl
     return forward_call(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File &quot;/repo/moe_benchmarks/megablocks_yamoe/.uvnote/cells/megablocks_run.py&quot;, line 81, in forward
     output, dummy_routing_weights = self.model(hidden_states)
                                     ^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/torch/nn/modules/module.py&quot;, line 1773, in _wrapped_call_impl
     return self._call_impl(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/torch/nn/modules/module.py&quot;, line 1784, in _call_impl
     return forward_call(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/layers.py&quot;, line 896, in forward
     output, expert_weights_out, *_ = moe_forward(
                                      ^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/layers.py&quot;, line 730, in moe_forward
     x, tokens_per_expert = forward_fn(**forward_args)
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/layers.py&quot;, line 457, in forward_once
     x = permute_and_compute(
         ^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/layers.py&quot;, line 401, in permute_and_compute
     x = ops.binned_gather(x, indices, bins, expert_capacity, top_k)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/torch/autograd/function.py&quot;, line 576, in apply
     return super().apply(*args, **kwargs)  # type: ignore[misc]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/ops/stk_autocast.py&quot;, line 30, in decorate_fwd
     return fwd(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/ops/binned_gather.py&quot;, line 26, in forward
     return kernels.binned_gather(x, indices, None, bins, bin_size, top_k)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/huggingface/hub/models--kernels-community--megablocks/snapshots/e0fb1437de3f8d7079c4da13be8cb64dc0cfcdd5/build/torch28-cxx11-cu128-x86_64-linux/megablocks/backend/kernels.py&quot;, line 419, in binned_gather
     _binned_copy[(num_experts, expert_capacity)](
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/jit.py&quot;, line 390, in &lt;lambda&gt;
     return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/autotuner.py&quot;, line 239, in run
     benchmark()
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/autotuner.py&quot;, line 228, in benchmark
     timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/autotuner.py&quot;, line 228, in &lt;dictcomp&gt;
     timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/autotuner.py&quot;, line 160, in _bench
     return self.do_bench(kernel_call, quantiles=(0.5, 0.2, 0.8))
            ^^^^^^^^^^^^^
   File &quot;/usr/lib/python3.11/functools.py&quot;, line 1001, in __get__
     val = self.func(instance)
           ^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/autotuner.py&quot;, line 121, in do_bench
     return driver.active.get_benchmarker()
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/driver.py&quot;, line 30, in __getattr__
     return getattr(self._initialize_obj(), name)
                    ^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/driver.py&quot;, line 26, in _initialize_obj
     self._obj = self._init_fn()
                 ^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/driver.py&quot;, line 12, in _create_driver
     return active_drivers[0]()
            ^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/backends/nvidia/driver.py&quot;, line 715, in __init__
     self.utils = CudaUtils()  # TODO: make static
                  ^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/backends/nvidia/driver.py&quot;, line 62, in __init__
     mod = compile_module_from_src(
           ^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/build.py&quot;, line 88, in compile_module_from_src
     so = _build(name, src_path, tmpdir, library_dirs or [], include_dirs or [], libraries or [])
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File &quot;/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/runtime/build.py&quot;, line 51, in _build
     subprocess.check_call(cc_cmd, stdout=subprocess.DEVNULL)
   File &quot;/usr/lib/python3.11/subprocess.py&quot;, line 413, in check_call
     raise CalledProcessError(retcode, cmd)
+subprocess.CalledProcessError: Command &#x27;[&#x27;/usr/bin/gcc&#x27;, &#x27;/tmp/tmptrubhjfl/cuda_utils.c&#x27;, &#x27;-O3&#x27;, &#x27;-shared&#x27;, &#x27;-fPIC&#x27;, &#x27;-Wno-psabi&#x27;, &#x27;-o&#x27;, &#x27;/tmp/tmptrubhjfl/cuda_utils.cpython-311-x86_64-linux-gnu.so&#x27;, &#x27;-lcuda&#x27;, &#x27;-L/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/backends/nvidia/lib&#x27;, &#x27;-L/usr/lib/x86_64-linux-gnu&#x27;, &#x27;-I/tmp/uvnote-run-68wjowzh/home/.cache/uv/environments-v2/megablocks-run-8802ebf6d3566120/lib/python3.11/site-packages/triton/backends/nvidia/include&#x27;, &#x27;-I/tmp/tmptrubhjfl&#x27;, &#x27;-I/usr/include/python3.11&#x27;]&#x27; returned non-zero exit status 1.</div>
 </div>
 </div>